From 15efd003340cca467afe1e239695d48d10ee8ad2 Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Wed, 9 Sep 2026 18:51:08 +0300 Subject: [PATCH 01/14] Upgrade GR00T integration to N1.7 DROID --- README.md | 2 +- docker/Dockerfile.groot | 10 +- docker/Makefile | 4 +- docker/README.md | 21 +- docker/docker-compose.phail.yml | 3 + docker/docker-compose.spoons-ablation.yml | 13 +- docker/docker-compose.yml | 9 +- docs/codecs.md | 34 +- docs/inference.md | 4 +- docs/model-selection.md | 22 +- docs/training-workflow.md | 11 +- positronic/cfg/ds/__init__.py | 2 +- positronic/offboard/README.md | 2 +- positronic/vendors/gr00t/README.md | 215 ++++-------- positronic/vendors/gr00t/__init__.py | 70 +--- positronic/vendors/gr00t/codecs.py | 306 +++++++----------- positronic/vendors/gr00t/server.py | 182 +++++------ positronic/vendors/gr00t/tests/test_codecs.py | 56 ++-- .../vendors/gr00t/tests/test_observation.py | 283 ++++------------ positronic/vendors/gr00t/tests/test_server.py | 33 ++ positronic/vendors/gr00t/tests/test_train.py | 33 ++ positronic/vendors/gr00t/train.py | 34 +- pyproject.toml | 1 + uv.lock | 16 + workflows/nebius/README.md | 4 +- workflows/nebius/convert.sh | 2 +- workflows/nebius/e2e.sh | 7 +- workflows/nebius/serve.sh | 2 +- 28 files changed, 541 insertions(+), 840 deletions(-) create mode 100644 positronic/vendors/gr00t/tests/test_train.py diff --git a/README.md b/README.md index 070edd85a..635d2453c 100644 --- a/README.md +++ b/README.md @@ -82,7 +82,7 @@ Positronic supports state-of-the-art foundation models with first-class workflow | Model | Capability | Training | Inference | Best For | |-------|-----------|----------|-----------|----------| | **[OpenPI (π₀.₅)](positronic/vendors/openpi/README.md)** | Most capable, generalist | Capable GPU (~78GB, LoRA) | Capable GPU (~62GB) | Complex multi-task manipulation | -| **[GR00T](positronic/vendors/gr00t/README.md)** | Generalist robot policy | Capable GPU (~50GB) | Smaller GPU (~7.5GB) | Logistics and industry applications | +| **[GR00T N1.7 DROID](positronic/vendors/gr00t/README.md)** | DROID robot policy | CUDA GPU | CUDA GPU | Joint control with 2 or 3 camera views | | **[LeRobot SmolVLA](positronic/vendors/lerobot/README.md)** | VLM-based, multi-task | Consumer GPU | Consumer GPU | Multi-task manipulation with language | | **[LeRobot ACT](positronic/vendors/lerobot_0_3_3/README.md)** | Single-task, efficient | Consumer GPU | Consumer GPU | Specific manipulation tasks | diff --git a/docker/Dockerfile.groot b/docker/Dockerfile.groot index 79bc9b406..ee9a6a315 100644 --- a/docker/Dockerfile.groot +++ b/docker/Dockerfile.groot @@ -1,14 +1,22 @@ -ARG BASE_IMAGE=positro/gr00t-base:latest +ARG BASE_IMAGE=positro/gr00t-base:240627d FROM ${BASE_IMAGE} +COPY --from=ghcr.io/astral-sh/uv:latest /uv /uvx /usr/local/bin/ + # Set working directory WORKDIR /positronic +RUN apt-get update && apt-get install -y --no-install-recommends libturbojpeg && rm -rf /var/lib/apt/lists/* + # Copy the positronic repository COPY . /positronic # Set Python path to include source root and gr00t ENV PYTHONPATH=/positronic:/gr00t +# Keep Positronic's dependencies separate from the checkpoint runtime. +ENV UV_PROJECT_ENVIRONMENT=/positronic/.venv VIRTUAL_ENV=/positronic/.venv +RUN --mount=type=cache,target=/root/.cache/uv uv sync --locked --python 3.12 --no-dev + # Default command CMD ["bash"] diff --git a/docker/Makefile b/docker/Makefile index d8840cab7..b4c082d0f 100644 --- a/docker/Makefile +++ b/docker/Makefile @@ -41,6 +41,8 @@ TAG_DREAMZERO_BASE_LATEST := $(IMAGE_NAME_DREAMZERO_BASE):latest # OpenPI base image (pulled from registry) TAG_OPENPI_BASE := $(IMAGE_NAME_OPENPI_BASE):latest +GROOT_BASE_IMAGE ?= positro/gr00t-base:240627d + help: @echo "Positronic Docker Build System - Makefile" @echo "" @@ -81,7 +83,7 @@ build-openpi: build-groot: build-training @echo "Building $(IMAGE_NAME_GROOT)..." - docker build --pull --platform linux/amd64 --build-arg BASE_IMAGE=positro/gr00t-base:latest -f Dockerfile.groot -t $(IMAGE_NAME_GROOT):local .. + docker build --platform linux/amd64 --build-arg BASE_IMAGE=$(GROOT_BASE_IMAGE) -f Dockerfile.groot -t $(IMAGE_NAME_GROOT):local .. build-dreamzero-base: @echo "Building $(IMAGE_NAME_DREAMZERO_BASE)..." diff --git a/docker/README.md b/docker/README.md index d74e05903..a43e0ab9f 100644 --- a/docker/README.md +++ b/docker/README.md @@ -40,7 +40,22 @@ Now you are ready to build our Docker. docker/build.sh ``` -## GR00T containers: uv mount caveat +## GR00T N1.7 containers -If you customize `docker-compose.yml` volumes, **do not bind-mount** your host `~/.local/share/uv` into `/root/.local/share/uv` for `positro/gr00t` images. -GR00T's `/.venv/bin/python` can be a symlink into the image's own uv-managed CPython under `/root/.local/share/uv/python/...`, and the bind mount can hide that target and cause `/.venv/bin/python` to fail with `ENOENT`. +`make build-groot` uses `GROOT_BASE_IMAGE` from the GR00T fork. The fork image contains +CUDA 12.8 and the upstream Python 3.12 environment at `/opt/gr00t-venv`. +Positronic installs its own locked environment at `/positronic/.venv`. +Both training and serving launch GR00T in its separate environment. + +Build the base in the [fork](https://github.com/Positronic-Robotics/gr00t), then build the adapter: + +```bash +# In the GR00T fork: +make -C docker build +# In Positronic: +make -C docker build-groot GROOT_BASE_IMAGE=positro/gr00t-base:local +IMAGE_TAG=local docker compose -f docker/docker-compose.yml run --rm --service-ports groot-server droid +``` + +Do not mount host uv interpreter directories over the image's interpreter directories. +See [GR00T](../positronic/vendors/gr00t/README.md) for conversion, fine-tuning and inference. diff --git a/docker/docker-compose.phail.yml b/docker/docker-compose.phail.yml index c7c375cd1..e97fd39ed 100644 --- a/docker/docker-compose.phail.yml +++ b/docker/docker-compose.phail.yml @@ -13,6 +13,9 @@ services: extends: file: docker-compose.yml service: groot-server + # These experiment checkpoints use the N1.6 action schema. + image: ${GR00T_N16_IMAGE:?Set GR00T_N16_IMAGE to a pinned GR00T N1.6 image} + entrypoint: ["uv", "run", "--no-sync", "--python", "3.13", "python", "-m", "positronic.vendors.gr00t.server"] container_name: phail-groot-server pull_policy: always command: ["phail"] diff --git a/docker/docker-compose.spoons-ablation.yml b/docker/docker-compose.spoons-ablation.yml index 040056a74..ec0fb2573 100644 --- a/docker/docker-compose.spoons-ablation.yml +++ b/docker/docker-compose.spoons-ablation.yml @@ -1,14 +1,17 @@ # GR00T spoons data ablation servers (~4.6GB each, 2 fit on desktop, 1 on notebook) # # Desktop (2 servers): -# IMAGE_TAG=latest CACHE_ROOT=/home/ docker --context desktop compose -f docker-compose.spoons-ablation.yml up spoons-100 spoons-50 +# GR00T_N16_IMAGE=positro/gr00t: CACHE_ROOT=/home/ docker --context desktop compose -f docker-compose.spoons-ablation.yml up spoons-100 spoons-50 # Notebook (1 server): -# IMAGE_TAG=latest docker --context notebook compose -f docker-compose.spoons-ablation.yml up spoons-25 +# GR00T_N16_IMAGE=positro/gr00t: docker --context notebook compose -f docker-compose.spoons-ablation.yml up spoons-25 services: spoons-100: extends: file: docker-compose.yml service: groot-server + # These experiment checkpoints use the N1.6 action schema. + image: ${GR00T_N16_IMAGE:?Set GR00T_N16_IMAGE to a pinned GR00T N1.6 image} + entrypoint: ["uv", "run", "--no-sync", "--python", "3.13", "python", "-m", "positronic.vendors.gr00t.server"] container_name: spoons-100 pull_policy: always command: @@ -21,6 +24,9 @@ services: extends: file: docker-compose.yml service: groot-server + # These experiment checkpoints use the N1.6 action schema. + image: ${GR00T_N16_IMAGE:?Set GR00T_N16_IMAGE to a pinned GR00T N1.6 image} + entrypoint: ["uv", "run", "--no-sync", "--python", "3.13", "python", "-m", "positronic.vendors.gr00t.server"] container_name: spoons-50 pull_policy: always command: @@ -33,6 +39,9 @@ services: extends: file: docker-compose.yml service: groot-server + # These experiment checkpoints use the N1.6 action schema. + image: ${GR00T_N16_IMAGE:?Set GR00T_N16_IMAGE to a pinned GR00T N1.6 image} + entrypoint: ["uv", "run", "--no-sync", "--python", "3.13", "python", "-m", "positronic.vendors.gr00t.server"] container_name: spoons-25 pull_policy: always command: diff --git a/docker/docker-compose.yml b/docker/docker-compose.yml index a02fdb4ff..d9f6b9b81 100644 --- a/docker/docker-compose.yml +++ b/docker/docker-compose.yml @@ -181,10 +181,7 @@ services: volumes: - ${CACHE_ROOT:-${HOME}}/.cache:/root/.cache - ${CACHE_ROOT:-${HOME}}/.aws:/root/.aws:ro - # NOTE: Do NOT mount `/root/.local/share/uv` for gr00t images. - # gr00t's `/.venv/bin/python` is a symlink into the image's own uv-managed CPython under - # `/root/.local/share/uv/python/...`. Bind-mounting the host uv dir can hide that target and - # make `/.venv/bin/python` fail with ENOENT. + # Keep image-owned uv interpreters visible; do not mount /root/.local/share/uv here. # Set HOME inside container to match the mount expectations environment: @@ -201,12 +198,12 @@ services: container_name: groot-train shm_size: 8g ipc: host # Enable shared memory access for training - entrypoint: ["uv", "run", "--python", "3.13", "python", "-m", "positronic.vendors.gr00t.train"] + entrypoint: ["uv", "run", "--no-sync", "--python", "3.12", "python", "-m", "positronic.vendors.gr00t.train"] groot-server: &groot-server-common <<: *groot-common container_name: groot-server - entrypoint: ["uv", "run", "--python", "3.13", "python", "-m", "positronic.vendors.gr00t.server"] + entrypoint: ["uv", "run", "--no-sync", "--python", "3.12", "python", "-m", "positronic.vendors.gr00t.server"] ports: - "8000:8000" diff --git a/docs/codecs.md b/docs/codecs.md index 92e021dce..3aef6fd98 100644 --- a/docs/codecs.md +++ b/docs/codecs.md @@ -118,31 +118,15 @@ cd docker && docker compose run --rm lerobot-convert convert \ See [`positronic/vendors/gr00t/codecs.py`](../positronic/vendors/gr00t/codecs.py). -| Codec | Observation | Action | Modality Configs | -|-------|-------------|--------|------------------| -| `ee_quat` | EE pose (quat) + grip + images (224x224) | Absolute EE position (quat) + grip | `ee`, `ee_rel` | -| `ee_rot6d` | EE pose (rot6d) + grip + images | Absolute EE position (rot6d) + grip | `ee_rot6d`, `ee_rot6d_rel` | -| `ee_quat_joints` | EE pose + joints + grip + images | Absolute EE position + grip | `ee_q`, `ee_q_rel` | -| `ee_rot6d_joints` | EE pose (rot6d) + joints + grip + images | Absolute EE position (rot6d) + grip | `ee_rot6d_q`, `ee_rot6d_q_rel` | -| `ee_quat_traj` | EE pose (quat) + grip + images | Absolute EE trajectory (quat) + grip (binarized) | `ee`, `ee_rel` | -| `ee_rot6d_traj` | EE pose (rot6d) + grip + images | Absolute EE trajectory (rot6d) + grip (binarized) | `ee_rot6d`, `ee_rot6d_rel` | -| `ee_quat_joints_traj` | EE pose + joints + grip + images | Absolute EE trajectory + grip (binarized) | `ee_q`, `ee_q_rel` | -| `ee_rot6d_joints_traj` | EE pose (rot6d) + joints + grip + images | Absolute EE trajectory (rot6d) + grip (binarized) | `ee_rot6d_q`, `ee_rot6d_q_rel` | -| `joints_traj` | Joints + grip + images (no EE pose) | Absolute joint trajectory + grip (binarized) | — | - -The codec must match the modality config during training. - -```bash -# Convert with codec -cd docker && docker compose run --rm lerobot-0_3_3-convert convert \ - --dataset.codec=@positronic.vendors.gr00t.codecs.ee_rot6d_joints \ - --output_dir=~/datasets/groot/my_task - -# Train with matching modality -cd docker && docker compose run --rm groot-train \ - --modality_config=ee_rot6d_q \ - --input_path=~/datasets/groot/my_task -``` +| Codec | Cameras | State and training actions | Inference actions | +|-------|---------|----------------------------|-------------------| +| `droid` | Exterior + wrist | Absolute EEF pose (XYZ + row-based rot6d), gripper, 7 joints | Absolute joint targets + binary gripper | +| `droid_three_cameras` | Two exteriors + wrist | Same as `droid` | Same as `droid` | + +Images use the upstream DROID client's 320×180 padded resize, then the checkpoint's native +preprocessing. GR00T owns pose-relative and joint-relative conversion. Training labels are +recorded state trajectories. Use the same camera layout for conversion and inference; +the published DROID checkpoint uses two cameras. See the [Docker workflow](../positronic/vendors/gr00t/README.md). ### OpenPI diff --git a/docs/inference.md b/docs/inference.md index 71cc4e9f6..c8b1ea714 100644 --- a/docs/inference.md +++ b/docs/inference.md @@ -6,7 +6,7 @@ Deploy trained policies for evaluation and production use. Positronic supports l Positronic's unified WebSocket protocol connects any hardware to any model (LeRobot, GR00T, OpenPI). The key benefit is running heavy models on powerful GPU hardware (OpenPI needs ~62GB, GR00T ~8GB) separate from the robot/simulator machine. -Each server carries a full **policy pipeline** — one chain naming the rig-side stack, the `remote` split marker, the server-side codec, and the model source that loads checkpoints (see `positronic.policy.spec`). The server runs the half right of the marker and declares the half left of it in its handshake; the client builds the declared stack automatically. Vendors ship their pipelines by name, and every name is a server subcommand — `groot-server ee_rot6d_joints` launches that one. The available names are listed in each vendor's README. +Each server carries a full **policy pipeline** — one chain naming the rig-side stack, the `remote` split marker, the server-side codec, and the model source that loads checkpoints (see `positronic.policy.spec`). The server runs the half right of the marker and declares the half left of it in its handshake; the client builds the declared stack automatically. Vendors ship their pipelines by name, and every name is a server subcommand — `groot-server droid` launches that one. The available names are listed in each vendor's README. **Start inference server:** ```bash @@ -20,7 +20,7 @@ cd docker && docker compose run --rm --service-ports lerobot-0_3_3-server ee \ --pipeline.source.checkpoints_dir=~/checkpoints/lerobot/experiment_v1/ # GR00T -cd docker && docker compose run --rm --service-ports groot-server ee_rot6d_joints \ +cd docker && docker compose run --rm --service-ports groot-server droid \ --pipeline.source.checkpoints_dir=~/checkpoints/groot/experiment_v1/ # OpenPI (--pipeline.ee_frame states the EE frame the checkpoint speaks; None means the rig's `default`) diff --git a/docs/model-selection.md b/docs/model-selection.md index 08b68c34a..fba59043f 100644 --- a/docs/model-selection.md +++ b/docs/model-selection.md @@ -11,9 +11,9 @@ Positronic supports three foundation models with different capabilities and reso | Aspect | OpenPI (π₀.₅) | GR00T | SmolVLA | LeRobot ACT | |--------|---------------|-------|---------|-------------| | **Capability** | Most capable, generalist | Generalist | Vision-language-action | Single-task specialist | -| **Training Hardware** | capable cloud GPU (~78GB, LoRA) | capable cloud GPU (~50GB) | Consumer GPU (RTX 3090, 4090) | Consumer GPU (RTX 3090, 4090) | +| **Training Hardware** | capable cloud GPU (~78GB, LoRA) | CUDA GPU (measure for batch/layout) | Consumer GPU (RTX 3090, 4090) | Consumer GPU (RTX 3090, 4090) | | **Training Time** | Multiple days | 0.5-2 days | Several hours | Several hours | -| **Inference Hardware** | GPU (~62GB, likely cloud) | GPU (~7.5GB, can run on robot) | Consumer GPU (4GB+) | Consumer GPU (4GB+) | +| **Inference Hardware** | GPU (~62GB, likely cloud) | CUDA GPU (measure for camera layout) | Consumer GPU (4GB+) | Consumer GPU (4GB+) | | **Inference Speed** | Moderate | Moderate | Moderate | Fast | | **Best For** | Complex multi-task manipulation, generalization | General robotics tasks | Language-conditioned manipulation | Specific manipulation tasks, fast iteration | | **When to Use** | Need generalization, multi-task scenarios, leveraging foundation models | Prefer NVIDIA stack | Language instructions, VLM backbone | Single task, resource constraints, rapid experimentation | @@ -36,20 +36,14 @@ Positronic supports three foundation models with different capabilities and reso → [OpenPI Documentation](../positronic/vendors/openpi/README.md) -### GR00T — NVIDIA's Generalist Robot Policy +### GR00T N1.7 DROID -**What it is:** NVIDIA's foundation model for generalist robot control. - -**Strengths:** -- Generalist capabilities -- Can run on smaller GPU (~7.5GB inference, can run closer to robot) -- Requires ~50GB for training (less than OpenPI) -- Training takes 1-2 days (faster than OpenPI) - -**Limitations:** -- Requires capable GPU for training -- Slower than single-task models +NVIDIA's robot policy, using the published DROID checkpoint. Supports wrist plus one exterior +camera by default, and wrist plus two exteriors through fine-tuning. It predicts joint-position +actions and preserves the checkpoint's image and relative-action processing. +Training and inference require a CUDA GPU. Measure memory requirements with the intended batch +size and camera layout. → [GR00T Documentation](../positronic/vendors/gr00t/README.md) diff --git a/docs/training-workflow.md b/docs/training-workflow.md index 62752f7dc..e9bbede82 100644 --- a/docs/training-workflow.md +++ b/docs/training-workflow.md @@ -69,7 +69,7 @@ See the [Codecs Guide](codecs.md) for detailed codec documentation. |-------|---------------| | **SmolVLA / LeRobot 0.4.x** | `ee`, `joints` (512x512 images) | | **LeRobot ACT (0.3.3)** | `ee`, `joints`, `ee_traj`, `joints_traj` | -| **GR00T** | `ee_rot6d_joints`, `ee_quat`, `ee_quat_joints` | +| **GR00T N1.7** | `droid`, `droid_three_cameras` | | **OpenPI** | `ee`, `ee_joints`, `droid` | ### S3 Support @@ -127,11 +127,10 @@ cd docker && docker compose run --rm lerobot-train full_finetune \ cd docker && docker compose run --rm groot-train \ --input_path=~/datasets/groot/stack_cubes \ --output_path=~/checkpoints/groot \ - --exp_name=experiment_v1 \ - --modality_config=ee_rot6d_q + --exp_name=experiment_v1 ``` -**Modality config must match codec** (see [GR00T README](../positronic/vendors/gr00t/README.md#1-prepare-data)). +The checkpoint supplies the model and action configuration. Camera names come from the converted dataset (see [GR00T README](../positronic/vendors/gr00t/README.md)). ### OpenPI Training @@ -197,7 +196,7 @@ cd docker && docker compose run --rm --service-ports lerobot-0_3_3-server ee \ **GR00T Server (naming the pipeline as the subcommand):** ```bash -cd docker && docker compose run --rm --service-ports groot-server ee_rot6d_joints \ +cd docker && docker compose run --rm --service-ports groot-server droid \ --pipeline.source.checkpoints_dir=~/checkpoints/groot/experiment_v1/ ``` @@ -287,7 +286,7 @@ cd docker && docker compose run --rm lerobot-convert convert \ # ACT, GR00T, OpenPI — use lerobot-0_3_3-convert cd docker && for pair in \ "lerobot_0_3_3.codecs.ee ~/datasets/lerobot_act/my_task" \ - "gr00t.codecs.ee_rot6d_joints ~/datasets/groot/my_task" \ + "gr00t.codecs.droid ~/datasets/groot/my_task" \ "openpi.codecs.ee ~/datasets/openpi/my_task"; do set -- $pair docker compose run --rm lerobot-0_3_3-convert convert \ diff --git a/positronic/cfg/ds/__init__.py b/positronic/cfg/ds/__init__.py index 436f2bd52..c6aa252cd 100644 --- a/positronic/cfg/ds/__init__.py +++ b/positronic/cfg/ds/__init__.py @@ -60,7 +60,7 @@ def apply_codec(dataset: Dataset, codec): positronic-to-lerobot convert \\ --dataset=@positronic.cfg.ds.apply_codec \\ --dataset.dataset=.internal.droid \\ - --dataset.codec=@positronic.vendors.gr00t.codecs.ee_quat \\ + --dataset.codec=@positronic.vendors.gr00t.codecs.droid \\ --output_dir=/data/lerobot_dataset """ return TransformedDataset(dataset, codec.training_encoder) diff --git a/positronic/offboard/README.md b/positronic/offboard/README.md index ddf7140e1..5699245c6 100644 --- a/positronic/offboard/README.md +++ b/positronic/offboard/README.md @@ -189,7 +189,7 @@ cd docker && docker compose run --rm --service-ports lerobot-server ee \ --pipeline.source.checkpoints_dir=~/checkpoints/lerobot/exp_v1 # GR00T server (swap hardware code stays the same) -cd docker && docker compose run --rm --service-ports groot-server ee_rot6d_joints \ +cd docker && docker compose run --rm --service-ports groot-server droid \ --pipeline.source.checkpoints_dir=~/checkpoints/groot/exp_v1 # Client connects the same way diff --git a/positronic/vendors/gr00t/README.md b/positronic/vendors/gr00t/README.md index 8276a81d7..c2510ac7c 100644 --- a/positronic/vendors/gr00t/README.md +++ b/positronic/vendors/gr00t/README.md @@ -1,175 +1,100 @@ -# GR00T in Positronic +# GR00T N1.7 DROID -## What is GR00T? +Positronic uses [`nvidia/GR00T-N1.7-DROID`](https://huggingface.co/nvidia/GR00T-N1.7-DROID) +through our [GR00T fork](https://github.com/Positronic-Robotics/gr00t). +The checkpoint defines the model architecture, image processor and relative-action conversion. +The adapter follows [upstream's DROID robot client](https://github.com/NVIDIA/Isaac-GR00T/tree/main/examples/DROID). -GR00T is [NVIDIA's](https://developer.nvidia.com/isaac/groot) generalist robot foundation model for versatile robot control. +## Representation -Positronic provides first-class support for GR00T including: -- Training on single capable server GPU (~50GB) -- Inference on smaller GPU (~7.5GB, can run closer to robot) -- Relative modalities support (uses same codecs, different groot model internally to match OpenPI's relative actions by default) -- Unified inference API compatible with all Positronic hardware -- Integration with our fork: [Positronic-Robotics/gr00t](https://github.com/Positronic-Robotics/gr00t), kept up to date with upstream +- `droid`: wrist + one exterior camera, matching the published checkpoint. +- `droid_three_cameras`: wrist + two exterior cameras. Fine-tune with this layout, then serve that checkpoint. +- RGB images first use the client's bilinear padded resize to **320×180**. The checkpoint processor + resizes the shortest edge to **256**, crops **95%**, resizes the shortest edge again, then runs + its vision processor. Positronic does not apply an additional square crop. +- State is absolute tool position + row-based 6D rotation, gripper position and seven joint positions. + Poses move from Positronic's default tool frame to `DROID_EE_FRAME`, then receive upstream's + DROID rotation correction. Gripper convention is **1 = closed**. +- Conversion writes recorded state trajectories as absolute action labels. The checkpoint processor + computes pose-relative EEF actions and joint offsets for training, then restores absolute actions + at inference. Positronic does not subtract poses or rotations itself. +- Inference executes the first **15 of 40** predicted joint targets at **15 Hz**, with the DROID + impedance settings and a gripper threshold of **0.5**, matching upstream's default execution horizon. -See [Model Selection Guide](../../docs/model-selection.md) for comparison with other options. +The published two-camera checkpoint does not consume an extra exterior view. Select +`droid_three_cameras` for both conversion and serving when fine-tuning with three views. +N1.6 checkpoints require an N1.6 image; their custom action schemas are incompatible with this adapter. -## Hardware Requirements +## Docker -| Phase | Requirement | Notes | -|-------|-------------|-------| -| **Training** | capable sever GPU (~50GB) | NVIDIA's training config optimized for a single capable GPU | -| **Inference** | GPU (~7.5GB) | RTX 4070, A10, or better (can run on robot) | -| **Training Time** | 0.5-2 days | Typical for GR00T | - -## Quick Start +Build the fork's base image, then the Positronic image: ```bash -# 1. Convert dataset (output_dir supports both local paths and s3://) -cd docker && docker compose run --rm lerobot-0_3_3-convert convert \ - --dataset.dataset.path=~/datasets/my_task_raw \ - --dataset.codec=@positronic.vendors.gr00t.codecs.ee_rot6d_joints \ - --output_dir=~/datasets/groot/my_task - -# 2. Train -cd docker && docker compose run --rm groot-train \ - --input_path=~/datasets/groot/my_task \ - --output_path=~/checkpoints/groot \ - --exp_name=my_task_v1 \ - --modality_config=ee_rot6d_q - -# 3. Serve -cd docker && docker compose run --rm --service-ports groot-server ee_rot6d_joints \ - --pipeline.source.checkpoints_dir=~/checkpoints/groot/my_task_v1/ - -# 4. Run inference -uv run --locked positronic eval run --eval=.sim.positronic.stack_cubes \ - --policy=.remote \ - --policy.url=localhost:8000 +# In the GR00T fork +make -C docker build +# In Positronic +make -C docker build-groot GROOT_BASE_IMAGE=positro/gr00t-base:local +cd docker +export IMAGE_TAG=local ``` -See [Training Workflow](../../docs/training-workflow.md) for detailed step-by-step instructions. - -## Available Codecs +The GR00T environment is `/opt/gr00t-venv` (Python 3.12, upstream locked dependencies). +Positronic has a separate environment at `/positronic/.venv`. Training and serving require a CUDA GPU. -GR00T supports multiple codecs with different rotation representations and observation spaces. +## Convert and fine-tune -| Codec | Observation | Action | Modality Config | Use Case | -|-------|-------------|--------|-----------------|----------| -| `ee_quat` | EE pose (quat) + grip + images | Absolute EE position (quat) + grip | `ee` | Default EE control, quaternion rotation | -| `ee_rot6d` | EE pose (rot6d) + grip + images | Absolute EE position (rot6d) + grip | `ee_rot6d` | 6D rotation representation | -| `ee_quat_joints` | EE pose + joints + grip + images | Absolute EE position + grip | `ee_q` | Combined EE + joint feedback | -| `ee_rot6d_joints` | EE pose (rot6d) + joints + grip + images | Absolute EE position (rot6d) + grip | `ee_rot6d_q` | 6D rotation + joint feedback (recommended) | +From Positronic's `docker` directory: -**Key features:** -- **Rotation representations**: Quaternion (4D) vs rot6d (6D continuous) -- **Joint feedback**: Optional joint position observations for richer state representation -- Images automatically resized to 224x224 -- Sets `gr00t_modality` metadata for training compatibility - -**Codec must match modality config during training:** - -| Codec | Training Modality | -|-------|-------------------| -| `ee_quat` | `ee` | -| `ee_rot6d` | `ee_rot6d` | -| `ee_quat_joints` | `ee_q` | -| `ee_rot6d_joints` | `ee_rot6d_q` | - -**Recommendation:** Use `ee_rot6d_joints` for best performance (6D rotation is continuous, joint feedback improves learning). - -See [Codecs Guide](../../docs/codecs.md) for comprehensive codec documentation. - -## Configuration Reference - -### Training Configuration +```bash +docker compose run --rm --pull never lerobot-0_3_3-convert convert \ + --dataset.codec=@positronic.vendors.gr00t.codecs.droid \ + --output_dir=~/datasets/groot/my_task -**Common parameters:** +docker compose run --rm groot-train \ + --input_path=~/datasets/groot/my_task \ + --output_path=~/checkpoints/groot \ + --exp_name=my_task \ + --num_train_steps=10000 +``` -| Parameter | Description | Default | Example | -|-----------|-------------|---------|---------| -| `--modality_config` | Modality configuration (must match codec) | `ee` | `ee_rot6d_q` | -| `--exp_name` | Experiment name (unique ID) | Required | `my_task_v1` | -| `--num_train_steps` | Total training steps | Config default | `100000` | -| `--learning_rate` | Override learning rate | Config default | `1e-4` | -| `--save_steps` | Checkpoint save interval | Config default | `10000` | -| `--num_workers` | Dataloader workers | Config default | `8` | -| `--resume` | Resume from existing checkpoint | `False` | `True` | -| `--output_path` | Checkpoint destination | Required | `~/checkpoints/groot` | +Supply the conversion command's dataset configuration for your recordings as usual. +For three views, replace the codec with `positronic.vendors.gr00t.codecs.droid_three_cameras`. +The launcher reads camera keys from `meta/modality.json`; no separate modality selection is needed. -**WandB logging:** Enabled by default if `WANDB_API_KEY` is set in `docker/.env.wandb`. +`--base_model` defaults to `nvidia/GR00T-N1.7-DROID`. Standard controls are `--batch_size`, +`--learning_rate`, `--num_train_steps`, `--save_steps`, `--num_workers` and `--resume=True`. +Resume restores the latest saved training state in the experiment directory. +The fork retains checkpoint architecture and preprocessing while applying upstream's standard +fine-tuning settings. Dataset statistics are computed by GR00T. -### Inference Server Configuration +## Serve -Every named policy pipeline is a server subcommand, pairing the codec with the matching modality -config: +Published checkpoint, without fine-tuning: ```bash -cd docker && docker compose run --rm --service-ports groot-server ee_rot6d_joints \ - --pipeline.source.checkpoints_dir=~/checkpoints/groot/my_task_v1/ \ - --port=8000 +docker compose run --rm --service-ports groot-server droid ``` -**Available pipelines** (the subcommand selects one): -- `ee` - End-effector pose (quaternion) -- `ee_joints` - End-effector pose + joint positions (quaternion) -- `ee_rot6d` - End-effector pose (rot6d) -- `ee_rot6d_joints` - End-effector pose + joint positions (rot6d, recommended) -- `ee_rot6d_rel` - End-effector pose (rot6d, relative actions) -- `ee_rot6d_joints_rel` - End-effector pose + joint positions (rot6d, relative actions) - -`serve` is `ee`. The `phail` and `sim_stack` subcommands are the same pipelines with their -`checkpoints_dir`/`recording_dir` bound. - -**Server parameters:** - -| Parameter | Description | Default | Example | -|-----------|-------------|---------|---------| -| subcommand | Named pipeline | `ee` | `ee_rot6d_joints` | -| `--pipeline.source.checkpoints_dir` | Experiment directory (contains `checkpoint-N` folders) | Required | `~/checkpoints/groot/my_task_v1/` | -| `--pipeline.source.checkpoint` | Specific checkpoint ID | Latest | `10000`, `50000` | -| `--port` | Server port | `8000` | `8001` | -| `--pipeline.source.modality_config` | Override the pipeline's paired modality config | Paired | `ee_rot6d_q` | - -**Session parameters:** a client can tune the served pipeline per connection via query params on the -session URL — e.g. `--policy.url='vm-h100:8000?codec.fps=10'` on the eval CLI. Values -must be JSON literals. The model source (`checkpoints_dir`, `--checkpoint`, modality config) is -fixed at launch and cannot be changed per session. - -## Troubleshooting +Fine-tuned checkpoint: -### GR00T Modality Mismatch +```bash +docker compose run --rm --service-ports groot-server droid \ + --pipeline.source.checkpoints_dir=~/checkpoints/groot/my_task +``` -**Problem:** Training or inference fails with modality-related errors +Select `droid_three_cameras` for a checkpoint trained on three views. +Use `--pipeline.source.checkpoint=10000` to select a saved step. Omit it to serve the latest. +A Hugging Face source uses `--pipeline.source.checkpoints_dir=hf://owner/model`. -**Cause:** Codec and modality config don't match +## Adapter parity tests -**Solution:** Use the correct pairing (see table in [Available Codecs](#available-codecs)): +The GR00T source is included in the image at `/gr00t`. From the image's `/positronic` directory: ```bash -# Codec: ee_rot6d_joints → Modality: ee_rot6d_q - -# Training -cd docker && docker compose run --rm groot-train \ - --modality_config=ee_rot6d_q \ - --input_path=~/datasets/groot/my_task # (converted with ee_rot6d_joints codec) - -# Inference (use the matching pipeline) -cd docker && docker compose run --rm --service-ports groot-server ee_rot6d_joints \ - --pipeline.source.checkpoints_dir=~/checkpoints/groot/my_task_v1/ +GR00T_REFERENCE_ROOT=/gr00t uv run --no-sync --python 3.12 pytest \ + -o addopts= positronic/vendors/gr00t/tests ``` -## See Also - -**Positronic Documentation:** -- [Model Selection Guide](../../docs/model-selection.md) — When to use GR00T vs OpenPI vs LeRobot -- [Codecs Guide](../../docs/codecs.md) — Understanding observation/action encoding -- [Training Workflow](../../docs/training-workflow.md) — Unified training steps across all models -- [Inference Guide](../../docs/inference.md) — Deployment and evaluation patterns - -**Other Models:** -- [OpenPI (π₀.₅)](../openpi/README.md) — Recommended for most tasks, most capable foundation model -- [LeRobot ACT](../lerobot/README.md) — Single-task transformer, fast training - -**External:** -- [NVIDIA GR00T](https://developer.nvidia.com/isaac/groot) — Official GR00T page -- [Positronic GR00T Fork](https://github.com/Positronic-Robotics/gr00t) — Our integration repository +The cross-repository tests compare encoded tool poses and image pixels directly against the +upstream DROID functions. Model forward and fine-tuning checks additionally require the checkpoint +weights and a GPU. diff --git a/positronic/vendors/gr00t/__init__.py b/positronic/vendors/gr00t/__init__.py index b6910125e..43e8977d4 100644 --- a/positronic/vendors/gr00t/__init__.py +++ b/positronic/vendors/gr00t/__init__.py @@ -1,65 +1,19 @@ -"""What positronic states about GR00T: the nested observation it takes, and the modality configs it serves.""" +"""GR00T DROID model and wire vocabulary.""" -from dataclasses import dataclass -from pathlib import Path - -# Keys of the nested observation a GR00T session takes, which the codec writes and a warmup rebuilds. VIDEO = 'video' STATE = 'state' LANGUAGE = 'language' - -WRIST_IMAGE = 'wrist_image' -EXTERIOR_IMAGE = 'exterior_image_1' - -# The state/action fields GR00T's data config declares and its model emits. ``GRIP`` shares a value with -# ``keys.GRIP`` by vocabulary, not by contract — renaming the positronic wire key must not rename the modality. -GRIP = 'grip' -EE_POSE = 'ee_pose' +WRIST_IMAGE = 'wrist_image_left' +EXTERIOR_IMAGE = 'exterior_image_1_left' +EXTERIOR_IMAGE_2 = 'exterior_image_2_left' +GRIP = 'gripper_position' +EE_POSE = 'eef_9d' JOINT_POSITION = 'joint_position' - TASK = 'annotation.language.language_instruction' +EMBODIMENT = 'oxe_droid_relative_eef_relative_joint' +BASE_MODEL = 'nvidia/GR00T-N1.7-DROID' +VENV = '/opt/gr00t-venv' -# The frame geometry GR00T is served at, as ``(width, height)``: what the rig is bounded to, what the codec -# resizes to, and what a warmup fills. -IMAGE_SIZE = (224, 224) - - -@dataclass(frozen=True) -class ModalityConfig: - """One GR00T modality config: the fork module registering it, and the observation it declares. - - ``path`` is relative to the gr00t checkout, which is where the subprocess runs. The rest is positronic's - own statement of what a checkpoint served under this config takes. The fork's config module is the other - statement, and only gr00t's venv can import it, so nothing reconciles the two but the test pairing each - config with the codec that feeds it. - - ``cameras`` and ``task_key`` default to what every config shipped here declares; gr00t's other embodiments - name their cameras and their language field differently, so a config of your own states its own. - """ - - path: Path - state: dict[str, int] - cameras: tuple[str, ...] = (WRIST_IMAGE, EXTERIOR_IMAGE) - task_key: str = TASK - - -_CONFIG_DIR = Path('gr00t/configs/data') - -_EE_QUAT = {GRIP: 1, EE_POSE: 7} -_EE_QUAT_JOINTS = {GRIP: 1, EE_POSE: 7, JOINT_POSITION: 7} -_EE_ROT6D = {GRIP: 1, EE_POSE: 9} -_EE_ROT6D_JOINTS = {GRIP: 1, EE_POSE: 9, JOINT_POSITION: 7} - -# A ``_rel`` config differs from its twin in the action space it trains, not in the observation it takes. -MODALITY_CONFIGS = { - # 7D xyz+quat configs (absolute actions) - 'ee': ModalityConfig(_CONFIG_DIR / 'positronic_ee.py', _EE_QUAT), - 'ee_q': ModalityConfig(_CONFIG_DIR / 'positronic_ee_joints.py', _EE_QUAT_JOINTS), - # 9D xyz+rot6d configs (supports both absolute and relative actions) - 'ee_rot6d': ModalityConfig(_CONFIG_DIR / 'positronic_ee_rot6d.py', _EE_ROT6D), - 'ee_rot6d_rel': ModalityConfig(_CONFIG_DIR / 'positronic_ee_rot6d_rel.py', _EE_ROT6D), - 'ee_rot6d_q': ModalityConfig(_CONFIG_DIR / 'positronic_ee_rot6d_joints.py', _EE_ROT6D_JOINTS), - 'ee_rot6d_q_rel': ModalityConfig(_CONFIG_DIR / 'positronic_ee_rot6d_joints_rel.py', _EE_ROT6D_JOINTS), - # Joint-space action config (for IK-derived targets) - 'joints': ModalityConfig(_CONFIG_DIR / 'positronic_joints.py', {GRIP: 1, JOINT_POSITION: 7}), -} +# Width, height at the DROID robot-client boundary; the checkpoint processor owns subsequent resizing/cropping. +IMAGE_SIZE = (320, 180) +STATE_DIMS = {EE_POSE: 9, GRIP: 1, JOINT_POSITION: 7} diff --git a/positronic/vendors/gr00t/codecs.py b/positronic/vendors/gr00t/codecs.py index 531e90cb0..8ab076c3e 100644 --- a/positronic/vendors/gr00t/codecs.py +++ b/positronic/vendors/gr00t/codecs.py @@ -1,230 +1,154 @@ -"""GR00T codecs: implementation classes and configuronic configs in one file.""" +"""DROID observations and joint-position actions for GR00T.""" from functools import partial -from typing import Any +from operator import itemgetter import configuronic as cfn import numpy as np -from PIL import Image as PilImage +from PIL import Image from positronic import geom, keys -from positronic.cfg import codecs -from positronic.dataset import transforms +from positronic.cfg.hardware.roboarm import DROID_IMPEDANCE from positronic.dataset import transforms as tf from positronic.dataset.episode import Episode -from positronic.dataset.signal import Signal from positronic.dataset.transforms import image -from positronic.dataset.transforms.episode import Derive, Get, Identity -from positronic.policy.codec import Codec, lerobot_image, lerobot_state +from positronic.dataset.transforms.episode import Derive +from positronic.drivers.roboarm import command, models +from positronic.policy.codec import ( + ActionHorizon, + ActionTimestamp, + BinarizeGripInference, + ChangeEEFrame, + Codec, + lerobot_action, + lerobot_image, + lerobot_state, +) from positronic.vendors import gr00t -RotRep = geom.Rotation.Representation - -class GrootObservationCodec(Codec): - """GR00T N1.6 observation encoder. +class DroidCodec(Codec): + """Encode poses in the DROID tool frame and decode upstream's absolute joint targets. - For training (training_encoder): derives flat keys for each state component. - For inference: encode() produces nested GR00T format (video/state/language). + Training uses recorded pose/joint/gripper trajectories as absolute action labels. GR00T's + checkpoint processor converts those labels to relative actions and back during inference. """ - def __init__( - self, - rotation_rep: RotRep | None = None, - include_joints: bool = False, - include_ee_pose: bool = True, - image_size: tuple[int, int] = gr00t.IMAGE_SIZE, - exterior_camera: str = keys.EXTERIOR_IMAGE, - wrist_camera: str = keys.WRIST_IMAGE, - num_joints: int = 7, - ): - self._rotation_rep = rotation_rep - self._include_joints = include_joints - self._include_ee_pose = include_ee_pose - self._image_size = image_size - self._exterior_camera = exterior_camera - self._wrist_camera = wrist_camera - self._num_joints = num_joints - - self._derive_transforms: dict[str, Any] = { - gr00t.GRIP: self._derive_grip, - gr00t.WRIST_IMAGE: partial(self._derive_image, wrist_camera), - gr00t.EXTERIOR_IMAGE: partial(self._derive_image, exterior_camera), - 'task': Get(keys.TASK, ''), - } - - state_meta: dict[str, Any] = {gr00t.GRIP: {'start': 0, 'end': 1, 'original_key': gr00t.GRIP}} - lerobot_features: dict[str, Any] = { - gr00t.GRIP: lerobot_state(1), - gr00t.WRIST_IMAGE: lerobot_image(*image_size), - gr00t.EXTERIOR_IMAGE: lerobot_image(*image_size), - } - - if include_ee_pose: - obs_ee_dim = rotation_rep.size + 3 if rotation_rep else 7 - state_meta[gr00t.EE_POSE] = {'start': 0, 'end': obs_ee_dim, 'original_key': gr00t.EE_POSE} - lerobot_features[gr00t.EE_POSE] = lerobot_state(obs_ee_dim) - self._derive_transforms[gr00t.EE_POSE] = self._derive_ee_pose - if include_joints: - state_meta[gr00t.JOINT_POSITION] = {'start': 0, 'end': num_joints, 'original_key': gr00t.JOINT_POSITION} - lerobot_features[gr00t.JOINT_POSITION] = lerobot_state(num_joints) - self._derive_transforms[gr00t.JOINT_POSITION] = self._derive_joints - - self._training_meta = { - 'gr00t_modality': { - gr00t.STATE: state_meta, - gr00t.VIDEO: { - gr00t.EXTERIOR_IMAGE: {'original_key': gr00t.EXTERIOR_IMAGE}, - gr00t.WRIST_IMAGE: {'original_key': gr00t.WRIST_IMAGE}, - }, - 'annotation': {'language.language_instruction': {'original_key': 'task_index'}}, - }, - 'lerobot_features': lerobot_features, - } - - def _derive_ee_pose(self, episode: Episode) -> Signal[Any]: - pose = episode[keys.EE_POSE] - if self._rotation_rep is not None: - pose = tf.recode_transform(RotRep.QUAT, self._rotation_rep, pose) - return tf.astype(pose, np.float32) - - def _derive_grip(self, episode: Episode) -> Signal[Any]: - def _reshape_to_1d(values): - arr = np.asarray(values, dtype=np.float32) - return arr.reshape(-1, 1) - - return transforms.Elementwise(episode[keys.GRIP], _reshape_to_1d) + # Matches GR00T's gr00t/data/state_action/droid_frame.py; row-based rot6d follows this correction. + _ROTATION_CORRECTION = np.array([[0, 0, -1], [-1, 0, 0], [0, 1, 0]], dtype=np.float64) - def _derive_joints(self, episode: Episode) -> Signal[Any]: - return tf.astype(episode[keys.JOINTS], np.float32) + def __init__(self, image_mappings: dict[str, str]): + self.image_mappings = dict(image_mappings) - def _derive_image(self, input_key: str, episode: Episode) -> Signal[Any]: - w, h = self._image_size - return image.resize_with_pad(w, h, signal=episode[input_key]) + @classmethod + def _encode_pose(cls, value): + pose = geom.Transform3D.from_vector(np.asarray(value), geom.Rotation.Representation.QUAT) + rotation = pose.rotation.as_rotation_matrix @ cls._ROTATION_CORRECTION + return np.concatenate([pose.translation, rotation[:2].reshape(6)]).astype(np.float32) - def _encode_ee_pose(self, inputs: dict[str, Any]) -> np.ndarray: - pose = np.asarray(inputs[keys.EE_POSE], dtype=np.float32).reshape(-1) - if self._rotation_rep is not None: - pose = geom.Transform3D.from_vector(pose, RotRep.QUAT).as_vector(self._rotation_rep).astype(np.float32) - return pose - - def _encode_image(self, input_key: str, inputs: dict[str, Any]) -> np.ndarray: - frame = inputs[input_key] - if not isinstance(frame, np.ndarray): - frame = np.asarray(frame) - w, h = self._image_size - return image.resize_with_pad_per_frame(w, h, PilImage.Resampling.BILINEAR, frame) - - def _decode_single(self, data: dict) -> dict: - return {} - - def encode(self, inputs: dict[str, Any]) -> dict[str, Any]: - grip = np.asarray(inputs[keys.GRIP], dtype=np.float32).reshape(-1) - state_dict: dict[str, Any] = {gr00t.GRIP: grip[np.newaxis, np.newaxis, ...]} - - if self._include_ee_pose: - ee_pose = self._encode_ee_pose(inputs) - state_dict[gr00t.EE_POSE] = ee_pose[np.newaxis, np.newaxis, ...] - if self._include_joints: - joints = np.asarray(inputs[keys.JOINTS], dtype=np.float32).reshape(-1) - state_dict[gr00t.JOINT_POSITION] = joints[np.newaxis, np.newaxis, ...] + @staticmethod + def _encode_image(frame): + return image.resize_with_pad_per_frame(*gr00t.IMAGE_SIZE, Image.Resampling.BILINEAR, np.asarray(frame)) + def encode(self, inputs: dict) -> dict: + state = { + gr00t.EE_POSE: self._encode_pose(inputs[keys.EE_POSE]), + gr00t.GRIP: np.asarray(inputs[keys.GRIP], dtype=np.float32).reshape(1), + gr00t.JOINT_POSITION: np.asarray(inputs[keys.JOINTS], dtype=np.float32).reshape(7), + } return { gr00t.VIDEO: { - gr00t.WRIST_IMAGE: self._encode_image(self._wrist_camera, inputs)[np.newaxis, np.newaxis, ...], - gr00t.EXTERIOR_IMAGE: self._encode_image(self._exterior_camera, inputs)[np.newaxis, np.newaxis, ...], + name: self._encode_image(inputs[source])[None, None] for name, source in self.image_mappings.items() }, - gr00t.STATE: state_dict, - gr00t.LANGUAGE: {gr00t.TASK: [[inputs.get(keys.TASK, '')]]}, + gr00t.STATE: {name: value[None, None] for name, value in state.items()}, + gr00t.LANGUAGE: {gr00t.TASK: [[inputs[keys.TASK]]]}, } - @property - def meta(self): - return {'image_sizes': self._image_size} - - @property - def training_encoder(self): - return Derive(meta=self._training_meta, **self._derive_transforms) - + def _decode_single(self, data: dict) -> dict: + return { + keys.ROBOT_COMMAND: command.JointPosition( + positions=np.asarray(data[gr00t.JOINT_POSITION]).reshape(7), mode=DROID_IMPEDANCE + ), + keys.TARGET_GRIP: np.asarray(data[gr00t.GRIP]).item(), + } -class _GrootActionModality(Codec): - """Bridges GR00T modality-keyed actions and flat action vectors. + def _derive_pose(self, episode: Episode): + return tf.Elementwise(episode[keys.EE_POSE], tf.lazy_sequence(self._encode_pose)) - Training: adds ``gr00t_modality.action`` metadata. - Inference decode: converts GR00T's ``{action_key: ..., 'grip': ...}`` output - into ``{'action': flat_vector}`` so the downstream action decoder can read it. - """ + @staticmethod + def _derive_grip(episode: Episode): + return tf.Elementwise(episode[keys.GRIP], lambda values: np.asarray(values, dtype=np.float32).reshape(-1, 1)) - def __init__(self, modality: dict[str, Any], action_key: str): - self._training_meta = {'gr00t_modality': {'action': modality}} - self._action_key = action_key + def _derive_actions(self, episode: Episode): + return tf.concat(self._derive_pose(episode), self._derive_grip(episode), episode[keys.JOINTS], dtype=np.float32) - def encode(self, data): - return data - - def _decode_single(self, data: dict) -> dict: - action_part = np.asarray(data[self._action_key], dtype=np.float32).reshape(-1) - grip_part = np.asarray(data[gr00t.GRIP], dtype=np.float32).reshape(-1) - return {'action': np.concatenate([action_part, grip_part])} + @staticmethod + def _derive_image(source: str, episode: Episode): + return image.resize_with_pad(*gr00t.IMAGE_SIZE, signal=episode[source]) @property def training_encoder(self): - return Identity(meta=self._training_meta) - - -@cfn.config(rotation_rep=None, include_joints=False, include_ee_pose=True, num_joints=7) -def groot_obs(rotation_rep: str | None, include_joints: bool, include_ee_pose: bool, num_joints: int): - """GR00T N1.6 observation encoder.""" - rot_rep = RotRep(rotation_rep) if rotation_rep else None - return GrootObservationCodec( - rotation_rep=rot_rep, include_joints=include_joints, include_ee_pose=include_ee_pose, num_joints=num_joints - ) + state_meta = {name: {'start': 0, 'end': dim, 'original_key': name} for name, dim in gr00t.STATE_DIMS.items()} + action_meta = {} + start = 0 + for name, dim in gr00t.STATE_DIMS.items(): + action_meta[name] = {'start': start, 'end': start + dim} + start += dim + meta = { + 'gr00t_modality': { + gr00t.STATE: state_meta, + 'action': action_meta, + gr00t.VIDEO: {name: {'original_key': name} for name in self.image_mappings}, + 'annotation': {gr00t.TASK.removeprefix('annotation.'): {'original_key': 'task_index'}}, + }, + 'lerobot_features': { + **{name: lerobot_state(dim) for name, dim in gr00t.STATE_DIMS.items()}, + **{name: lerobot_image(*gr00t.IMAGE_SIZE) for name in self.image_mappings}, + 'action': lerobot_action(start), + }, + } + return Derive( + meta=meta, + **{ + gr00t.EE_POSE: self._derive_pose, + gr00t.GRIP: self._derive_grip, + gr00t.JOINT_POSITION: lambda episode: tf.Elementwise( + episode[keys.JOINTS], partial(np.asarray, dtype=np.float32) + ), + 'task': itemgetter(keys.TASK), + 'action': self._derive_actions, + **{name: partial(self._derive_image, source) for name, source in self.image_mappings.items()}, + }, + ) + @property + def meta(self): + return {'image_sizes': dict.fromkeys(self.image_mappings.values(), gr00t.IMAGE_SIZE)} -@cfn.config(action_key=gr00t.EE_POSE, action_dim=7) -def groot_action(base, action_key: str, action_dim: int): - """Wrap an action codec with GR00T modality metadata and decode adapter. - Composition is ``base | _GrootActionModality`` so that on decode (right-to-left) - the modality adapter runs first, converting GR00T's modality-keyed output - into a flat ``action`` vector that ``base`` can decode. - """ - return base | _GrootActionModality( - {action_key: {'start': 0, 'end': action_dim}, gr00t.GRIP: {'start': action_dim, 'end': action_dim + 1}}, - action_key=action_key, +@cfn.config( + image_mappings={gr00t.EXTERIOR_IMAGE: keys.EXTERIOR_IMAGE, gr00t.WRIST_IMAGE: keys.WRIST_IMAGE}, + fps=15.0, + execution_horizon=15, + ee_frame=models.DROID_EE_FRAME, +) +def droid(image_mappings: dict[str, str], fps: float, execution_horizon: int, ee_frame: geom.Transform3D): + """DROID's image, frame, gripper and 15 Hz joint-control conventions.""" + if fps <= 0 or not 1 <= execution_horizon <= 40: + raise ValueError('fps must be positive and execution_horizon must be between 1 and 40') + return ( + ActionHorizon(execution_horizon / fps) + | ActionTimestamp(fps=fps) + | BinarizeGripInference() + | ChangeEEFrame(ee_frame) + | DroidCodec(image_mappings) ) -_ee_action = groot_action.override(base=codecs.absolute_pos_action) -_rot6d_obs = groot_obs.override(rotation_rep='rot6d') -_rot6d_action = _ee_action.override(**{'base.rotation_rep': 'rot6d', 'action_dim': 9}) - -ee_quat = codecs.compose.override(obs=groot_obs, action=_ee_action) -ee_quat_joints = ee_quat.override(**{'obs.include_joints': True}) -ee_rot6d = codecs.compose.override(obs=_rot6d_obs, action=_rot6d_action) -phail_v1 = ee_rot6d.override(action=codecs.phail_v1_execution.override(action=_rot6d_action)) -ee_rot6d_joints = ee_rot6d.override(**{'obs.include_joints': True}) - -_traj_action = _ee_action.override(base=codecs.traj_ee_action) -_rot6d_traj_action = _rot6d_action.override(base=codecs.traj_ee_action.override(rotation_rep='rot6d')) - -ee_quat_traj = codecs.compose.override(obs=groot_obs, action=_traj_action, binarize_grip=(keys.GRIP,)) -ee_rot6d_traj = codecs.compose.override(obs=_rot6d_obs, action=_rot6d_traj_action, binarize_grip=(keys.GRIP,)) -ee_quat_joints_traj = ee_quat_traj.override(**{'obs.include_joints': True}) -ee_rot6d_joints_traj = ee_rot6d_traj.override(**{'obs.include_joints': True}) - -joints_traj = codecs.compose.override( - obs=groot_obs.override(include_joints=True, include_ee_pose=False), - action=groot_action.override( - base=codecs.absolute_joints_action.override(tgt_joints_key=keys.JOINTS, tgt_grip_key=keys.GRIP), - action_key=gr00t.JOINT_POSITION, - ), - binarize_grip=(keys.GRIP,), -) - -# IK variants: GR00T obs (with joints) + IK joint-space action via groot_action wrapper -ee_joints_ik = codecs.compose.override( - obs=groot_obs.override(include_joints=True), - action=groot_action.override(base=codecs.ik_joints_action, action_key=gr00t.JOINT_POSITION), +droid_three_cameras = droid.override( + image_mappings={ + gr00t.EXTERIOR_IMAGE: keys.EXTERIOR_IMAGE, + gr00t.EXTERIOR_IMAGE_2: keys.EXTERIOR_IMAGE_2, + gr00t.WRIST_IMAGE: keys.WRIST_IMAGE, + } ) -ee_joints_ik_sim = ee_joints_ik.override(**{'action.base.solver': 'lm'}) diff --git a/positronic/vendors/gr00t/server.py b/positronic/vendors/gr00t/server.py index 2cbf2846c..50b049c46 100644 --- a/positronic/vendors/gr00t/server.py +++ b/positronic/vendors/gr00t/server.py @@ -1,4 +1,3 @@ -import io import logging import os import subprocess @@ -8,6 +7,7 @@ import configuronic as cfn import msgpack +import msgpack_numpy as mnp import numpy as np import pos3 import zmq @@ -18,7 +18,6 @@ from positronic.offboard.server_utils import run_with_progress, wait_for_subprocess_ready, warmup from positronic.policy import Policy, Session from positronic.policy import keys as policy_keys -from positronic.policy.codec import RestrictImageSize from positronic.policy.layers import ChunkedSchedule, StopOnFault from positronic.policy.spec import ModelSource, remote from positronic.utils.checkpoints import list_checkpoints @@ -28,14 +27,8 @@ logger = logging.getLogger(__name__) -########################################################################################### -# ZMQ client code for communicating with gr00t N1.6 server -# Adapted from gr00t/policy/server_client.py -########################################################################################### - - class MsgSerializer: - """Message serializer for ZMQ communication (N1.6 format).""" + """N1.7's msgpack-numpy wire format, excluding pickle-bearing object arrays.""" @staticmethod def to_bytes(data: Any) -> bytes: @@ -43,27 +36,26 @@ def to_bytes(data: Any) -> bytes: @staticmethod def from_bytes(data: bytes) -> Any: - return msgpack.unpackb(data, object_hook=MsgSerializer.decode_custom_classes) + return msgpack.unpackb(data, object_hook=MsgSerializer.decode_custom_classes, raw=False) @staticmethod def decode_custom_classes(obj): - if not isinstance(obj, dict): - return obj - if '__ndarray_class__' in obj: - return np.load(io.BytesIO(obj['as_npy']), allow_pickle=False) - return obj + if isinstance(obj, dict): + if obj.get(b'nd', obj.get('nd')) and obj.get(b'kind', obj.get('kind')) in (b'O', 'O'): + raise ValueError('Object arrays are not supported by the GR00T wire protocol') + if obj.get('__ModalityConfig__'): + return obj['as_json'] + return mnp.decode(obj) @staticmethod def encode_custom_classes(obj): - if isinstance(obj, np.ndarray): - output = io.BytesIO() - np.save(output, obj, allow_pickle=False) - return {'__ndarray_class__': True, 'as_npy': output.getvalue()} - return obj + if isinstance(obj, np.ndarray) and obj.dtype.hasobject: + raise TypeError('Object arrays are not supported by the GR00T wire protocol') + return mnp.encode(obj) class PolicyClient: - """Client for communicating with GR00T N1.6 PolicyServer via ZMQ.""" + """Client for communicating with GR00T N1.7 PolicyServer via ZMQ.""" def __init__(self, host: str = 'localhost', port: int = 5555, timeout_ms: int = 15000): self.context = zmq.Context() @@ -73,6 +65,8 @@ def __init__(self, host: str = 'localhost', port: int = 5555, timeout_ms: int = self._init_socket() def _init_socket(self): + if hasattr(self, 'socket'): + self.socket.close(linger=0) self.socket = self.context.socket(zmq.REQ) self.socket.setsockopt(zmq.RCVTIMEO, self.timeout_ms) self.socket.setsockopt(zmq.SNDTIMEO, self.timeout_ms) @@ -95,6 +89,7 @@ def call_endpoint(self, endpoint: str, data: dict | None = None, requires_input: self.socket.send(MsgSerializer.to_bytes(request)) message = self.socket.recv() except zmq.error.Again as err: + self._init_socket() raise RuntimeError( f'Timeout after {self.timeout_ms}ms calling endpoint "{endpoint}" at {self.host}:{self.port}' ) from err @@ -115,7 +110,7 @@ def reset(self) -> dict[str, Any]: return self.call_endpoint('reset', {'options': None}) def close(self): - self.socket.close() + self.socket.close(linger=0) self.context.term() @@ -127,16 +122,8 @@ def close(self): class Gr00tSubprocess: """Manages the gr00t ZMQ server subprocess.""" - def __init__( - self, - checkpoint_dir: str, - modality_config_path: Path, - groot_venv_path: str, - zmq_port: int = 5555, - ready_timeout: float = 120.0, - ): + def __init__(self, checkpoint_dir: str, groot_venv_path: str, zmq_port: int = 5555, ready_timeout: float = 120.0): self.checkpoint_dir = checkpoint_dir - self.modality_config_path = modality_config_path self.groot_venv_path = groot_venv_path self.zmq_port = zmq_port self.ready_timeout = ready_timeout @@ -149,8 +136,7 @@ def start(self, on_progress: Callable[[str], None] | None = None): command = [python_bin, 'gr00t/eval/run_gr00t_server.py'] command.extend(['--model_path', str(self.checkpoint_dir)]) - command.extend(['--embodiment_tag', 'NEW_EMBODIMENT']) - command.extend(['--modality_config_path', str(self.modality_config_path)]) + command.extend(['--embodiment-tag', gr00t.EMBODIMENT]) command.extend(['--host', '127.0.0.1']) command.extend(['--port', str(self.zmq_port)]) @@ -238,22 +224,29 @@ def _step_id(raw: str) -> str: return str(int(raw)) if raw.isdigit() else raw -def _warm_observation(modality: gr00t.ModalityConfig) -> dict[str, Any]: - """Zero-filled inputs in GR00T's nested format, carrying the state block ``modality`` declares. - - Leading axes are ``(batch, time)``, the way a session hands one step over. - """ +def _warm_observation(modalities: dict) -> dict[str, Any]: + """Build a valid current-frame DROID observation from the loaded checkpoint's modalities.""" + for name in (gr00t.VIDEO, gr00t.STATE): + if modalities[name]['delta_indices'] != [0]: + raise ValueError(f'DROID adapter requires current-frame {name}, got {modalities[name]}') width, height = gr00t.IMAGE_SIZE - frame = np.zeros((1, 1, height, width, 3), dtype=np.uint8) + state = { + name: np.zeros((1, 1, gr00t.STATE_DIMS[name]), dtype=np.float32) + for name in modalities[gr00t.STATE]['modality_keys'] + } + state[gr00t.EE_POSE][..., 3:] = [1, 0, 0, 0, 1, 0] return { - gr00t.VIDEO: dict.fromkeys(modality.cameras, frame), - gr00t.STATE: {key: np.zeros((1, 1, dim), dtype=np.float32) for key, dim in modality.state.items()}, - gr00t.LANGUAGE: {modality.task_key: [['']]}, + gr00t.VIDEO: { + name: np.zeros((1, 1, height, width, 3), dtype=np.uint8) + for name in modalities[gr00t.VIDEO]['modality_keys'] + }, + gr00t.STATE: state, + gr00t.LANGUAGE: {gr00t.TASK: [['pick up the object']]}, } class Gr00tSource(ModelSource): - """GR00T checkpoints under ``checkpoints_dir``, each served through a dedicated ZMQ subprocess. + """A Hugging Face model (``hf://owner/model``) or a directory of fine-tuned checkpoints. Model ids are checkpoint step numbers (``'5000'`` for ``checkpoint-5000``). ``load`` downloads the checkpoint and boots the gr00t subprocess; the returned policy owns the subprocess. @@ -261,26 +254,16 @@ class Gr00tSource(ModelSource): def __init__( self, - checkpoints_dir: str, + checkpoints_dir: str = 'hf://' + gr00t.BASE_MODEL, checkpoint: str | None = None, - modality_config: str | gr00t.ModalityConfig = 'ee', - groot_venv_path: str = '/.venv/', + groot_venv_path: str = gr00t.VENV, zmq_port: int = 5555, - ready_timeout: float = 120.0, + ready_timeout: float = 600.0, ): - if isinstance(modality_config, str): - if modality_config not in gr00t.MODALITY_CONFIGS: - raise ValueError( - f'Unknown modality config: {modality_config}. Available: {sorted(gr00t.MODALITY_CONFIGS)}. ' - 'A config of your own is passed as a ModalityConfig, which states the state block to warm it with' - ) - self._modality = gr00t.MODALITY_CONFIGS[modality_config] - else: - self._modality = modality_config self.checkpoints_dir = checkpoints_dir.rstrip('/') + if self._is_hub_model and checkpoint is not None: + raise ValueError('checkpoint step selection applies only to fine-tuned checkpoint directories') self.checkpoint = checkpoint - # What to call the config being served: the alias where there is one, else the module it points at. - self.modality_config = modality_config if isinstance(modality_config, str) else str(self._modality.path) self.groot_venv_path = groot_venv_path self.zmq_port = zmq_port self.ready_timeout = ready_timeout @@ -295,7 +278,13 @@ def _raw_for(self, model_id: str) -> str: return r raise ValueError(f'Checkpoint not found: {model_id}. Available: {self.get_models()}') + @property + def _is_hub_model(self) -> bool: + return self.checkpoints_dir.startswith('hf://') + def get_models(self) -> list[str]: + if self._is_hub_model: + return [self.checkpoints_dir.removeprefix('hf://')] return [_step_id(r) for r in self._raw_ids()] def resolve(self, model_id: str | None) -> str: @@ -304,6 +293,11 @@ def resolve(self, model_id: str | None) -> str: The zero-padding a directory may carry stays out of the public id; ``load`` puts it back to reach the directory. """ + if self._is_hub_model: + only_model = self.get_models()[0] + if model_id is not None and model_id != only_model: + raise ValueError(f'This source serves only {only_model}') + return only_model if model_id is None and self.checkpoint is not None: model_id = str(self.checkpoint).strip('/') if model_id is None: @@ -311,16 +305,18 @@ def resolve(self, model_id: str | None) -> str: return _step_id(self._raw_for(model_id)) def load(self, model_id: str, on_progress: Callable[[str], None] | None = None) -> Policy: - checkpoint_path = f'{self.checkpoints_dir}/checkpoint-{self._raw_for(model_id)}' - logger.info(f'Downloading checkpoint {checkpoint_path}') - checkpoint_dir = run_with_progress( - lambda: pos3.download(checkpoint_path, exclude=['optimizer.pt']), - f'Downloading checkpoint checkpoint-{model_id}', - on_progress, - ) + if self._is_hub_model: + self.resolve(model_id) + checkpoint_dir = self.checkpoints_dir + else: + checkpoint_path = f'{self.checkpoints_dir}/checkpoint-{self._raw_for(model_id)}' + checkpoint_dir = run_with_progress( + lambda: pos3.download(checkpoint_path, exclude=['optimizer.pt']), + f'Downloading checkpoint checkpoint-{model_id}', + on_progress, + ) groot = Gr00tSubprocess( checkpoint_dir=str(checkpoint_dir), - modality_config_path=self._modality.path, groot_venv_path=self.groot_venv_path, zmq_port=self.zmq_port, ready_timeout=self.ready_timeout, @@ -329,7 +325,8 @@ def load(self, model_id: str, on_progress: Callable[[str], None] | None = None) groot.start(on_progress) policy = Gr00tPolicy(groot, str(checkpoint_dir)) # The subprocess initializes CUDA on its first forward, which outlasts a rig's inference timeout. - warmup(policy, _warm_observation(self._modality), on_progress) + modalities = groot.client.call_endpoint('get_modality_config', requires_input=False) + warmup(policy, _warm_observation(modalities), on_progress) except Exception: groot.stop() raise @@ -338,7 +335,7 @@ def load(self, model_id: str, on_progress: Callable[[str], None] | None = None) def meta(self, model_id: str) -> dict[str, Any]: return { policy_keys.TYPE: 'groot', - 'modality_config': self.modality_config, + 'embodiment': gr00t.EMBODIMENT, policy_keys.EXPERIMENT_NAME: self.checkpoints_dir.split('/')[-1] or '', } @@ -351,52 +348,21 @@ def meta(self, model_id: str) -> dict[str, Any]: gr00t_source = cfn.Config(Gr00tSource) -# No ``ee_frame``: every checkpoint served here was trained on poses the rig reported in its ``default``, -# so none has a transform to declare. -@cfn.config(codec=codecs.ee_quat, source=gr00t_source) +@cfn.config(codec=codecs.droid, source=gr00t_source) def pipeline(codec, source): - return StopOnFault() | ChunkedSchedule() | RestrictImageSize(*gr00t.IMAGE_SIZE) | remote | codec | source - - -# Each entry pairs the codec with the matching GR00T modality config; they must agree with training. -ee = pipeline -ee_joints = pipeline.override(codec=codecs.ee_quat_joints, **{'source.modality_config': 'ee_q'}) -ee_rot6d = pipeline.override(codec=codecs.ee_rot6d, **{'source.modality_config': 'ee_rot6d'}) -ee_rot6d_joints = pipeline.override(codec=codecs.ee_rot6d_joints, **{'source.modality_config': 'ee_rot6d_q'}) -ee_rot6d_rel = pipeline.override(codec=codecs.ee_rot6d, **{'source.modality_config': 'ee_rot6d_rel'}) -ee_rot6d_joints_rel = pipeline.override(codec=codecs.ee_rot6d_joints, **{'source.modality_config': 'ee_rot6d_q_rel'}) -# The sim_stack checkpoint was trained on inverted-grip (1 = open) sim data, hence flip_grip. -sim_stack_pipe = pipeline.override( - codec=codecs.ee_rot6d.override(flip_grip=True), **{'source.modality_config': 'ee_rot6d'} -) + """Schedule DROID joint commands while the server codec performs checkpoint-specific conversion.""" + return StopOnFault() | ChunkedSchedule() | remote | codec | source -# Every pipeline is a subcommand, and so is every deployment — a pipeline with its checkpoints bound. +droid = pipeline +droid_three_cameras = pipeline.override(codec=codecs.droid_three_cameras) COMMANDS = { - 'serve': serve.override(pipeline=ee), - 'ee': serve.override(pipeline=ee), - 'ee_joints': serve.override(pipeline=ee_joints), - 'ee_rot6d': serve.override(pipeline=ee_rot6d), - 'ee_rot6d_joints': serve.override(pipeline=ee_rot6d_joints), - 'ee_rot6d_rel': serve.override(pipeline=ee_rot6d_rel), - 'ee_rot6d_joints_rel': serve.override(pipeline=ee_rot6d_joints_rel), - 'phail': serve.override( - pipeline=ee_rot6d_rel.override( - codec=codecs.phail_v1, - **{'source.checkpoints_dir': 's3://checkpoints/phail_unified/groot/270226-ee_rot6d_rel/'}, - ), - recording_dir='s3://inference/phail_unified/server_recordings/groot/270226-ee_rot6d_rel/', - ), - 'sim_stack': serve.override( - pipeline=sim_stack_pipe.override(**{ - 'source.checkpoints_dir': 's3://checkpoints/sim_stack/groot/ee_rot6d/230226/' - }), - recording_dir='s3://inference/sim_stack/server_recordings/groot/230226/', - ), + 'serve': serve.override(pipeline=droid), + 'droid': serve.override(pipeline=droid), + 'droid_three_cameras': serve.override(pipeline=droid_three_cameras), } if __name__ == '__main__': init_logging() - with pos3.mirror(): - cfn.cli(COMMANDS) + cfn.cli(COMMANDS) diff --git a/positronic/vendors/gr00t/tests/test_codecs.py b/positronic/vendors/gr00t/tests/test_codecs.py index cce2a364d..13b70c079 100644 --- a/positronic/vendors/gr00t/tests/test_codecs.py +++ b/positronic/vendors/gr00t/tests/test_codecs.py @@ -2,39 +2,23 @@ import pytest from positronic import keys -from positronic.geom import Rotation -from positronic.vendors.gr00t import EE_POSE, GRIP, JOINT_POSITION -from positronic.vendors.gr00t.codecs import ee_quat, joints_traj - -_T0_OBS = {keys.OBS_TIME_NS: 0} - - -def test_ee_quat_decodes_modality_keyed_actions(): - """GR00T models return modality-keyed dicts per action step, not a flat ``action`` vector; - the codec chain must convert this format into robot commands.""" - codec = ee_quat() - - ee_pose = np.concatenate([Rotation.identity.as_quat, [0.1, 0.2, 0.3]]).astype(np.float32) - model_output = [{EE_POSE: ee_pose, GRIP: np.float32(0.5)} for _ in range(3)] - - decoded = codec.decode(model_output) - assert len(decoded) == 4 # 3 actions + timestamp sentinel - for d in decoded[:-1]: - assert keys.ROBOT_COMMAND in d - assert keys.TARGET_GRIP in d - assert keys.ACTION_TIMESTAMP in d - assert decoded[-1] == {keys.ACTION_TIMESTAMP: pytest.approx(3 / 15.0)} # timestamp sentinel - - -def test_joints_traj_decodes_modality_keyed_actions(): - codec = joints_traj() - - joint_pos = np.array([0.1, -0.2, 0.3, 0.4, -0.5, 0.6, 0.7], dtype=np.float32) - model_output = [{JOINT_POSITION: joint_pos, GRIP: np.float32(0.8)} for _ in range(3)] - - decoded = codec.decode(model_output) - assert len(decoded) == 4 # 3 actions + timestamp sentinel - for d in decoded[:-1]: - assert keys.ROBOT_COMMAND in d - assert keys.TARGET_GRIP in d - assert decoded[-1] == {keys.ACTION_TIMESTAMP: pytest.approx(3 / 15.0)} # timestamp sentinel +from positronic.cfg.hardware.roboarm import DROID_IMPEDANCE +from positronic.vendors import gr00t +from positronic.vendors.gr00t.codecs import droid + + +def test_droid_executes_fifteen_joint_targets_at_15hz_and_binarizes_grip(): + codec = droid() + targets = np.arange(40 * 7, dtype=np.float32).reshape(40, 7) / 100 + output = [ + {gr00t.JOINT_POSITION: q, gr00t.GRIP: [0.5 if i % 2 else 0.51], gr00t.EE_POSE: np.zeros(9)} + for i, q in enumerate(targets) + ] + decoded = codec.decode(output) + assert len(decoded) == 16 + for i, item in enumerate(decoded[:-1]): + np.testing.assert_array_equal(item[keys.ROBOT_COMMAND].positions, targets[i]) + assert item[keys.ROBOT_COMMAND].mode == DROID_IMPEDANCE + assert item[keys.TARGET_GRIP] == (0.0 if i % 2 else 1.0) + assert item[keys.ACTION_TIMESTAMP] == pytest.approx(i / 15) + assert decoded[-1] == {keys.ACTION_TIMESTAMP: 1.0} diff --git a/positronic/vendors/gr00t/tests/test_observation.py b/positronic/vendors/gr00t/tests/test_observation.py index a70752766..ef8773b05 100644 --- a/positronic/vendors/gr00t/tests/test_observation.py +++ b/positronic/vendors/gr00t/tests/test_observation.py @@ -1,228 +1,83 @@ -"""Tests for GrootObservationCodec.""" - +import importlib.util +import os from pathlib import Path import numpy as np import pytest +from scipy.spatial.transform import Rotation from positronic import geom, keys -from positronic.vendors.gr00t import GRIP, JOINT_POSITION, LANGUAGE, MODALITY_CONFIGS, VIDEO, ModalityConfig, codecs -from positronic.vendors.gr00t.codecs import GrootObservationCodec -from positronic.vendors.gr00t.server import _warm_observation - -RotRep = geom.Rotation.Representation +from positronic.dataset.episode import EpisodeContainer +from positronic.dataset.tests.utils import DummySignal +from positronic.drivers.roboarm import models +from positronic.vendors import gr00t +from positronic.vendors.gr00t.codecs import droid, droid_three_cameras @pytest.fixture -def sample_inputs(): - """Sample raw inputs for inference encoding.""" +def observation(): + pose = geom.Transform3D([0.3, -0.2, 0.5], geom.Rotation.from_euler([0.4, -0.3, 0.7])) return { - keys.EE_POSE: np.array([0.1, 0.2, 0.3, 0.0, 0.0, 0.0, 1.0]), # xyz + quat (w,x,y,z) - keys.GRIP: np.array([0.5]), - keys.JOINTS: np.array([0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7]), - keys.WRIST_IMAGE: np.random.randint(0, 255, (480, 640, 3), dtype=np.uint8), - keys.EXTERIOR_IMAGE: np.random.randint(0, 255, (480, 640, 3), dtype=np.uint8), - keys.TASK: 'pick up the cube', + keys.EE_POSE: pose.as_vector(geom.Rotation.Representation.QUAT), + keys.GRIP: 0.25, + keys.JOINTS: np.arange(7, dtype=np.float64) / 10, + keys.WRIST_IMAGE: np.random.default_rng(1).integers(0, 256, (377, 611, 3), dtype=np.uint8), + keys.EXTERIOR_IMAGE: np.random.default_rng(2).integers(0, 256, (240, 320, 3), dtype=np.uint8), + keys.EXTERIOR_IMAGE_2: np.full((180, 320, 3), 73, dtype=np.uint8), + keys.TASK: 'Put the cup on the plate', } -class TestGrootObservationCodec: - """Tests for GrootObservationCodec training and inference modes.""" - - # --- Inference encoding tests --- - - def test_encode_basic(self, sample_inputs): - """Test basic inference encoding without rotation conversion.""" - codec = GrootObservationCodec(rotation_rep=None, include_joints=False) - result = codec.encode(sample_inputs) - - assert 'video' in result - assert 'state' in result - assert 'language' in result - - assert 'wrist_image' in result['video'] - assert 'exterior_image_1' in result['video'] - assert result['video']['wrist_image'].shape == (1, 1, 224, 224, 3) - assert result['video']['exterior_image_1'].shape == (1, 1, 224, 224, 3) - - assert 'ee_pose' in result['state'] - assert 'grip' in result['state'] - assert 'joint_position' not in result['state'] - assert result['state']['ee_pose'].shape == (1, 1, 7) - assert result['state']['grip'].shape == (1, 1, 1) - - assert result['language']['annotation.language.language_instruction'] == [['pick up the cube']] - - def test_encode_with_rot6d(self, sample_inputs): - """Test inference encoding with rot6d conversion.""" - codec = GrootObservationCodec(rotation_rep=RotRep.ROT6D, include_joints=False) - result = codec.encode(sample_inputs) - - assert result['state']['ee_pose'].shape == (1, 1, 9) - - ee_pose = result['state']['ee_pose'][0, 0] - assert np.allclose(ee_pose[:3], sample_inputs[keys.EE_POSE][:3]) - - expected_rot6d = geom.Rotation.from_quat(sample_inputs[keys.EE_POSE][3:7]).as_rot6d - assert np.allclose(ee_pose[3:], expected_rot6d, atol=1e-6) - - def test_encode_with_joints(self, sample_inputs): - """Test inference encoding with joint positions.""" - codec = GrootObservationCodec(rotation_rep=None, include_joints=True) - result = codec.encode(sample_inputs) - - assert 'joint_position' in result['state'] - assert result['state']['joint_position'].shape == (1, 1, 7) - assert np.allclose(result['state']['joint_position'][0, 0], sample_inputs[keys.JOINTS]) - - def test_encode_with_rot6d_and_joints(self, sample_inputs): - """Test inference encoding with both rot6d and joints.""" - codec = GrootObservationCodec(rotation_rep=RotRep.ROT6D, include_joints=True) - result = codec.encode(sample_inputs) - - assert result['state']['ee_pose'].shape == (1, 1, 9) - assert result['state']['grip'].shape == (1, 1, 1) - assert result['state']['joint_position'].shape == (1, 1, 7) - - def test_encode_missing_task(self, sample_inputs): - """Test inference encoding handles missing task gracefully.""" - del sample_inputs[keys.TASK] - codec = GrootObservationCodec() - result = codec.encode(sample_inputs) - - assert result['language']['annotation.language.language_instruction'] == [['']] - - # --- Rot6d conversion correctness tests --- - - def test_rot6d_identity_quaternion(self, sample_inputs): - """Test rot6d conversion with identity quaternion.""" - sample_inputs[keys.EE_POSE] = np.array([1.0, 2.0, 3.0, 1.0, 0.0, 0.0, 0.0]) - - codec = GrootObservationCodec(rotation_rep=RotRep.ROT6D) - result = codec.encode(sample_inputs) - - ee_pose = result['state']['ee_pose'][0, 0] - expected_rot6d = np.array([1.0, 0.0, 0.0, 0.0, 1.0, 0.0]) - assert np.allclose(ee_pose[3:], expected_rot6d, atol=1e-6) - - def test_rot6d_90deg_rotation(self, sample_inputs): - """Test rot6d conversion with 90 degree rotation around Z.""" - quat = np.array([0.0, 0.0, np.sin(np.pi / 4), np.cos(np.pi / 4)]) - sample_inputs[keys.EE_POSE] = np.array([1.0, 2.0, 3.0, *quat]) - - codec = GrootObservationCodec(rotation_rep=RotRep.ROT6D) - result = codec.encode(sample_inputs) - - ee_pose = result['state']['ee_pose'][0, 0] - expected_rot6d = geom.Rotation.from_quat(quat).as_rot6d - assert np.allclose(ee_pose[3:], expected_rot6d, atol=1e-6) - - # --- Output key tests --- - - def test_output_keys_basic(self): - """Test that codec outputs correct keys for training.""" - codec = GrootObservationCodec(rotation_rep=None, include_joints=False) - - expected = {'ee_pose', 'grip', 'wrist_image', 'exterior_image_1', 'task'} - assert set(codec._derive_transforms.keys()) == expected - - def test_output_keys_with_joints(self): - """Test that codec includes joint_position when enabled.""" - codec = GrootObservationCodec(rotation_rep=RotRep.ROT6D, include_joints=True) - - assert 'joint_position' in codec._derive_transforms - - # --- Metadata tests --- - - def test_training_meta(self): - """Test that training metadata is computed from constructor params.""" - codec = GrootObservationCodec(rotation_rep=RotRep.ROT6D, include_joints=True) - meta = codec._training_meta - - assert 'gr00t_modality' in meta - assert 'lerobot_features' in meta - assert 'joint_position' in meta['lerobot_features'] - assert meta['lerobot_features']['ee_pose']['shape'] == (9,) - - def test_training_meta_no_joints(self): - """Test that training metadata excludes joints when not enabled.""" - codec = GrootObservationCodec(rotation_rep=None, include_joints=False) - meta = codec._training_meta - - assert 'joint_position' not in meta['lerobot_features'] - assert meta['lerobot_features']['ee_pose']['shape'] == (7,) - - # --- Edge cases --- - - def test_non_square_input_image(self, sample_inputs): - """Test that non-square images are properly resized with padding.""" - sample_inputs[keys.WRIST_IMAGE] = np.random.randint(0, 255, (100, 200, 3), dtype=np.uint8) - - codec = GrootObservationCodec() - result = codec.encode(sample_inputs) - - assert result['video']['wrist_image'].shape == (1, 1, 224, 224, 3) - - def test_custom_image_size(self, sample_inputs): - """Test custom image size.""" - codec = GrootObservationCodec(image_size=(128, 128)) - result = codec.encode(sample_inputs) - - assert result['video']['wrist_image'].shape == (1, 1, 128, 128, 3) - assert result['video']['exterior_image_1'].shape == (1, 1, 128, 128, 3) - - def test_custom_camera_keys(self, sample_inputs): - """Test custom camera key mapping.""" - sample_inputs['cam1'] = sample_inputs.pop(keys.WRIST_IMAGE) - sample_inputs['cam2'] = sample_inputs.pop(keys.EXTERIOR_IMAGE) - - codec = GrootObservationCodec(wrist_camera='cam1', exterior_camera='cam2') - result = codec.encode(sample_inputs) - - assert result['video']['wrist_image'].shape == (1, 1, 224, 224, 3) - assert result['video']['exterior_image_1'].shape == (1, 1, 224, 224, 3) - - -def _shapes(observation: dict) -> dict: - """The nested observation reduced to the shape of every leaf, which is what a backend accepts or rejects.""" - return {name: {key: np.asarray(v).shape for key, v in block.items()} for name, block in observation.items()} - - -# Each pair is one deployment's codec and the GR00T modality config it was trained under. They must describe the -# same observation, and until the fork's own config module can be read from here, this is what says so. -# rules-allow: hardcoded-keys — the pairing is the assertion; reading it from the code under test would pass -# whatever that code held. -@pytest.mark.parametrize( - 'codec_name, modality_config', - [ - ('ee_quat', 'ee'), - ('ee_quat_joints', 'ee_q'), - ('ee_rot6d', 'ee_rot6d'), - ('ee_rot6d_joints', 'ee_rot6d_q'), - ('ee_rot6d', 'ee_rot6d_rel'), - ('ee_rot6d_joints', 'ee_rot6d_q_rel'), - ('joints_traj', 'joints'), - ], -) -def test_warmup_observation_matches_what_the_paired_codec_encodes(codec_name, modality_config, sample_inputs): - encoded = getattr(codecs, codec_name).instantiate().encode(sample_inputs) - - warm = _warm_observation(MODALITY_CONFIGS[modality_config]) - - assert _shapes(warm) == _shapes(encoded) - - -def test_a_custom_config_warms_at_the_cameras_and_language_field_it_declares(): - # Names GR00T's other embodiments use, which none of the configs shipped here declare. - camera, task_key = 'ego_view', 'annotation.human.coarse_action' - custom = ModalityConfig( - path=Path('gr00t/configs/data/my_own.py'), - state={GRIP: 1, JOINT_POSITION: 7}, - cameras=(camera,), - task_key=task_key, +@pytest.mark.parametrize('config', [droid, droid_three_cameras]) +def test_training_and_inference_encode_the_same_absolute_state_and_images(config, observation): + codec = config() + episode = EpisodeContainer({ + name: value if name == keys.TASK else DummySignal([0, 1], [value, value]) for name, value in observation.items() + }) + training = codec.training_encoder(episode) + encoded = codec.encode(observation) + for name, value in encoded[gr00t.STATE].items(): + assert np.asarray(training[name][0][0]).dtype == np.float32 + np.testing.assert_allclose(training[name][0][0], value[0, 0], atol=1e-6) + for name, frames in encoded[gr00t.VIDEO].items(): + assert frames.shape == (1, 1, 180, 320, 3) + np.testing.assert_array_equal(training[name][0][0], frames[0, 0]) + expected_action = np.concatenate([encoded[gr00t.STATE][name][0, 0] for name in gr00t.STATE_DIMS]) + np.testing.assert_allclose(training['action'][0][0], expected_action) + + +def test_three_camera_configuration_uses_a_distinct_second_external_image(observation): + encoded = droid_three_cameras().encode(observation) + assert len(encoded[gr00t.VIDEO]) == 3 + np.testing.assert_array_equal( + encoded[gr00t.VIDEO][gr00t.EXTERIOR_IMAGE_2][0, 0], observation[keys.EXTERIOR_IMAGE_2] ) - - warm = _warm_observation(custom) - - assert set(warm[VIDEO]) == {camera} - assert set(warm[LANGUAGE]) == {task_key} + del observation[keys.EXTERIOR_IMAGE_2] + with pytest.raises(KeyError): + droid_three_cameras().encode(observation) + + +def test_droid_frame_and_pixels_match_upstream_robot_client(observation): + reference = os.environ.get('GR00T_REFERENCE_ROOT') + if reference is None: + pytest.skip('Set GR00T_REFERENCE_ROOT to the GR00T checkout for cross-repository parity') + loaded = {} + for name, path in {'frame': 'gr00t/data/state_action/droid_frame.py', 'image': 'examples/DROID/utils.py'}.items(): + spec = importlib.util.spec_from_file_location(name, Path(reference) / path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + loaded[name] = module + raw_pose = geom.Transform3D.from_vector(observation[keys.EE_POSE], geom.Rotation.Representation.QUAT) + tool_pose = raw_pose * models.DROID_EE_FRAME + upstream_pose = np.concatenate([ + tool_pose.translation, + Rotation.from_matrix(tool_pose.rotation.as_rotation_matrix).as_euler('XYZ'), + ]) + encoded = droid().encode(observation) + np.testing.assert_allclose( + encoded[gr00t.STATE][gr00t.EE_POSE][0, 0], loaded['frame'].compute_eef_9d(upstream_pose), atol=1e-6 + ) + for name, source in {gr00t.EXTERIOR_IMAGE: keys.EXTERIOR_IMAGE, gr00t.WRIST_IMAGE: keys.WRIST_IMAGE}.items(): + expected = loaded['image'].resize_with_pad(observation[source], 180, 320) + np.testing.assert_array_equal(encoded[gr00t.VIDEO][name][0, 0], expected) diff --git a/positronic/vendors/gr00t/tests/test_server.py b/positronic/vendors/gr00t/tests/test_server.py index 2bea27e35..c0484308c 100644 --- a/positronic/vendors/gr00t/tests/test_server.py +++ b/positronic/vendors/gr00t/tests/test_server.py @@ -1,3 +1,9 @@ +import msgpack +import msgpack_numpy +import numpy as np +import pytest + +from positronic.vendors import gr00t from positronic.vendors.gr00t import server as gr00t_server @@ -16,3 +22,30 @@ def test_zero_padded_checkpoints_are_served_under_the_id_they_advertise(monkeypa assert source.resolve(None) == '10000' # The raw suffix survives only where it is needed — reaching the directory. assert source._raw_for('5000') == '005000' + + +def test_msgpack_numpy_preserves_actions_and_camera_arrays(): + + actions = {gr00t.JOINT_POSITION: np.arange(280, dtype=np.float32).reshape(1, 40, 7)} + upstream_bytes = msgpack.packb((actions, {}), default=msgpack_numpy.encode) + decoded, _ = gr00t_server.MsgSerializer.from_bytes(upstream_bytes) + np.testing.assert_array_equal(decoded[gr00t.JOINT_POSITION], actions[gr00t.JOINT_POSITION]) + image = np.arange(180 * 320 * 3, dtype=np.uint8).reshape(1, 1, 180, 320, 3) + encoded = gr00t_server.MsgSerializer.to_bytes({gr00t.VIDEO: image}) + np.testing.assert_array_equal(msgpack.unpackb(encoded, object_hook=msgpack_numpy.decode)[gr00t.VIDEO], image) + + +def test_serializer_rejects_pickle_bearing_arrays(): + + with pytest.raises(TypeError, match='Object arrays'): + gr00t_server.MsgSerializer.to_bytes(np.array([object()], dtype=object)) + for payload in ({b'nd': True, b'kind': b'O'}, {'nd': 1, 'kind': 'O'}): + with pytest.raises(ValueError, match='Object arrays'): + gr00t_server.MsgSerializer.from_bytes(msgpack.packb(payload)) + + +def test_published_checkpoint_is_served_without_a_local_checkpoint_scan(monkeypatch): + + source = gr00t_server.Gr00tSource() + assert source.get_models() == [gr00t.BASE_MODEL] + assert source.resolve(None) == gr00t.BASE_MODEL diff --git a/positronic/vendors/gr00t/tests/test_train.py b/positronic/vendors/gr00t/tests/test_train.py new file mode 100644 index 000000000..c010120ee --- /dev/null +++ b/positronic/vendors/gr00t/tests/test_train.py @@ -0,0 +1,33 @@ +import json +from contextlib import nullcontext +from unittest.mock import Mock + +import pytest + +from positronic.vendors import gr00t +from positronic.vendors.gr00t import train + + +@pytest.mark.parametrize('resume', [False, True]) +def test_finetuning_forwards_dataset_cameras_and_resume_to_gr00t(tmp_path, monkeypatch, resume): + dataset = tmp_path / 'dataset' + (dataset / 'meta').mkdir(parents=True) + cameras = [gr00t.EXTERIOR_IMAGE, gr00t.EXTERIOR_IMAGE_2, gr00t.WRIST_IMAGE] + (dataset / 'meta' / 'modality.json').write_text(json.dumps({gr00t.VIDEO: dict.fromkeys(cameras, {})})) + output = tmp_path / 'output' + output.mkdir() + sync = Mock(return_value=output) + run = Mock() + monkeypatch.setattr(train.pos3, 'mirror', nullcontext) + monkeypatch.setattr(train.pos3, 'download', lambda _: dataset) + monkeypatch.setattr(train.pos3, 'sync', sync) + monkeypatch.setattr(train.utils, 'save_run_metadata', Mock()) + monkeypatch.setattr(train.subprocess, 'run', run) + train.main(input_path=str(dataset), output_path=str(output), exp_name='droid', resume=resume, num_train_steps=2) + arguments = run.call_args.args[0] + assert arguments[arguments.index('--base-model-path') + 1] == gr00t.BASE_MODEL + offset = arguments.index('--video-keys') + 1 + assert arguments[offset : offset + 3] == cameras + assert ('--resume-from-checkpoint' in arguments) == resume + assert sync.call_args.kwargs['delete_remote'] == (not resume) + assert run.call_args.kwargs['check'] is True diff --git a/positronic/vendors/gr00t/train.py b/positronic/vendors/gr00t/train.py index 6809e71c6..27d4c5a97 100644 --- a/positronic/vendors/gr00t/train.py +++ b/positronic/vendors/gr00t/train.py @@ -1,3 +1,4 @@ +import json import os import subprocess from pathlib import Path @@ -6,7 +7,7 @@ import pos3 from positronic import utils -from positronic.vendors.gr00t import MODALITY_CONFIGS +from positronic.vendors import gr00t def cleanup_old_optimizers(output_dir: str, keep_last_n: int = 2): @@ -19,44 +20,43 @@ def cleanup_old_optimizers(output_dir: str, keep_last_n: int = 2): print(f'Deleted {opt_file}') -@cfn.config(num_train_steps=None, groot_venv_path='/.venv/', modality_config='ee') +@cfn.config(num_train_steps=None, groot_venv_path=gr00t.VENV, base_model=gr00t.BASE_MODEL, batch_size=64) def main( input_path: str, output_path: str, exp_name: str, - modality_config: str, + base_model: str, + batch_size: int, num_train_steps, groot_venv_path: str, - learning_rate: float = None, - save_steps: int = None, + learning_rate: float | None = None, + save_steps: int | None = None, resume: bool = False, - num_workers: int = None, + num_workers: int | None = None, keep_optimizers_for_last_n: int = 2, ): exp_name = str(exp_name) groot_root = Path(__file__).parents[4] / 'gr00t' python_bin = str(Path(groot_venv_path).expanduser() / 'bin' / 'python') - known = MODALITY_CONFIGS.get(modality_config) - modality_config_path = known.path if known is not None else Path(modality_config) with pos3.mirror(): dataset_local_path = pos3.download(input_path) + with (Path(dataset_local_path) / 'meta' / 'modality.json').open() as f: + video_keys = list(json.load(f)[gr00t.VIDEO]) output_path = output_path.rstrip('/') # When resuming, don't delete existing checkpoint files output_dir = pos3.sync(output_path + '/' + exp_name, delete_remote=not resume) prefix = 'resume_metadata' if resume else 'run_metadata' utils.save_run_metadata(output_dir, patterns=['*.py', '*.toml'], prefix=prefix) - # Calculate save_steps: 20 checkpoints per run, but at least every 2000 steps - if save_steps is None and num_train_steps is not None: - save_steps = max(num_train_steps // 20, 2000) - - # N1.6 uses launch_finetune.py with new CLI format (auto-resumes if checkpoint exists) command = [python_bin, 'gr00t/experiment/launch_finetune.py'] - command.extend(['--base_model_path', 'nvidia/GR00T-N1.6-3B']) + command.extend(['--base-model-path', base_model]) command.extend(['--dataset_path', str(dataset_local_path)]) - command.extend(['--modality_config_path', str(modality_config_path)]) - command.extend(['--embodiment_tag', 'NEW_EMBODIMENT']) + command.extend(['--video-keys', *video_keys]) + command.extend(['--embodiment-tag', gr00t.EMBODIMENT]) + command.extend(['--global-batch-size', str(batch_size)]) + if resume: + command.append('--resume-from-checkpoint') command.extend(['--output_dir', str(output_dir)]) command.extend(['--num_gpus', '1']) command.extend(['--save_total_limit', '9999']) # Keep all checkpoints @@ -71,7 +71,7 @@ def main( command.append('--use-wandb') env = os.environ.copy() - print(f'Running command: `{" ".join(command)}`\n with env: {env}') + print(f'Running command: {command}') subprocess.run(command, check=True, cwd=str(groot_root), env=env) # Clean up optimizer state from old checkpoints to save space diff --git a/pyproject.toml b/pyproject.toml index 281f7640e..404fa12e8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,6 +26,7 @@ dependencies = [ "fire", "httpx", "msgpack", + "msgpack-numpy==0.4.8", "mujoco", "numpy", "opencv-python-headless", diff --git a/uv.lock b/uv.lock index 8e162a212..4fb559769 100644 --- a/uv.lock +++ b/uv.lock @@ -3520,6 +3520,20 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/e5/db/0314e4e2db56ebcf450f277904ffd84a7988b9e5da8d0d61ab2d057df2b6/msgpack-1.1.2-cp313-cp313-win_arm64.whl", hash = "sha256:e69b39f8c0aa5ec24b57737ebee40be647035158f14ed4b40e6f150077e21a84", size = 64118, upload-time = "2025-10-08T09:15:23.402Z" }, ] +[[package]] +name = "msgpack-numpy" +version = "0.4.8" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "msgpack" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "(extra == 'extra-10-positronic-lerobot' and extra == 'extra-10-positronic-lerobot-0-3-3') or (extra == 'extra-10-positronic-lerobot' and extra == 'extra-10-positronic-molmoact2') or (extra == 'extra-10-positronic-lerobot-0-3-3' and extra == 'extra-10-positronic-molmoact2') or (extra != 'extra-10-positronic-molmoact2' and extra == 'extra-10-positronic-yam') or (extra != 'extra-10-positronic-lerobot' and extra != 'extra-10-positronic-lerobot-0-3-3' and extra == 'extra-10-positronic-yam')" }, + { name = "numpy", version = "2.3.5", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-10-positronic-lerobot' or extra == 'extra-10-positronic-lerobot-0-3-3' or extra == 'extra-10-positronic-molmoact2' or extra != 'extra-10-positronic-yam'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/08/94/61e8aee142733ebfdc400a05bdac6e1763c4514bba3b42743d223f388450/msgpack-numpy-0.4.8.tar.gz", hash = "sha256:c667d3180513422f9c7545be5eec5d296dcbb357e06f72ed39cc683797556e69", size = 10923, upload-time = "2022-06-09T03:43:08.739Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/9b/5d/f25ac7d4fb77cbd53ddc6d05d833c6bf52b12770a44fa9a447eed470ca9a/msgpack_numpy-0.4.8-py2.py3-none-any.whl", hash = "sha256:773c19d4dfbae1b3c7b791083e2caf66983bb19b40901646f61d8731554ae3da", size = 6919, upload-time = "2022-06-09T03:43:06.82Z" }, +] + [[package]] name = "msgspec" version = "0.21.1" @@ -4901,6 +4915,7 @@ dependencies = [ { name = "httpx" }, { name = "jinja2" }, { name = "msgpack" }, + { name = "msgpack-numpy" }, { name = "mujoco" }, { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "(extra == 'extra-10-positronic-lerobot' and extra == 'extra-10-positronic-lerobot-0-3-3') or (extra == 'extra-10-positronic-lerobot' and extra == 'extra-10-positronic-molmoact2') or (extra == 'extra-10-positronic-lerobot-0-3-3' and extra == 'extra-10-positronic-molmoact2') or (extra != 'extra-10-positronic-molmoact2' and extra == 'extra-10-positronic-yam') or (extra != 'extra-10-positronic-lerobot' and extra != 'extra-10-positronic-lerobot-0-3-3' and extra == 'extra-10-positronic-yam')" }, { name = "numpy", version = "2.3.5", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-10-positronic-lerobot' or extra == 'extra-10-positronic-lerobot-0-3-3' or extra == 'extra-10-positronic-molmoact2' or extra != 'extra-10-positronic-yam'" }, @@ -5017,6 +5032,7 @@ requires-dist = [ { name = "lerobot", extras = ["smolvla"], marker = "extra == 'lerobot'", specifier = "==0.4.3" }, { name = "linuxpy", marker = "extra == 'hardware'" }, { name = "msgpack" }, + { name = "msgpack-numpy", specifier = "==0.4.8" }, { name = "mujoco" }, { name = "numpy" }, { name = "nvidia-ml-py", marker = "extra == 'telemetry'" }, diff --git a/workflows/nebius/README.md b/workflows/nebius/README.md index 2abf2c448..8160bcd1e 100644 --- a/workflows/nebius/README.md +++ b/workflows/nebius/README.md @@ -144,7 +144,7 @@ bash workflows/nebius/convert.sh openpi \ bash workflows/nebius/convert.sh gr00t \ --dataset.dataset=@positronic.cfg.ds.sim.sim_stack_cubes \ - --dataset.codec=@positronic.vendors.gr00t.codecs.ee_rot6d_joints \ + --dataset.codec=@positronic.vendors.gr00t.codecs.droid \ --output_dir=s3:///sim_stack_cubes_gr00t/ ``` @@ -250,7 +250,7 @@ bash workflows/nebius/serve.sh openpi my-openpi ee \ --pipeline.source.checkpoints_dir=s3:///checkpoints/openpi// \ --pipeline.ee_frame=None -bash workflows/nebius/serve.sh gr00t groot-server ee_rot6d_rel \ +bash workflows/nebius/serve.sh gr00t groot-server droid \ --pipeline.source.checkpoints_dir=s3:///checkpoints/groot// ``` diff --git a/workflows/nebius/convert.sh b/workflows/nebius/convert.sh index 782cced93..0ca584e91 100644 --- a/workflows/nebius/convert.sh +++ b/workflows/nebius/convert.sh @@ -45,7 +45,7 @@ Examples: bash workflows/nebius/convert.sh gr00t \ --dataset.dataset=@positronic.cfg.ds.sim.sim_stack_cubes \ - --dataset.codec=@positronic.vendors.gr00t.codecs.ee_rot6d_joints \ + --dataset.codec=@positronic.vendors.gr00t.codecs.droid \ --output_dir=s3:///sim_stack_cubes_gr00t/ EOF exit 1 diff --git a/workflows/nebius/e2e.sh b/workflows/nebius/e2e.sh index 0c021175b..a32dc1082 100644 --- a/workflows/nebius/e2e.sh +++ b/workflows/nebius/e2e.sh @@ -48,7 +48,7 @@ case "$VENDOR" in lerobot_0_3_3) CODEC=positronic.vendors.lerobot_0_3_3.codecs.ee ;; lerobot) CODEC=positronic.vendors.lerobot.codecs.ee ;; openpi) CODEC=positronic.vendors.openpi.codecs.ee ;; - gr00t) CODEC=positronic.vendors.gr00t.codecs.ee_rot6d ;; + gr00t) CODEC=positronic.vendors.gr00t.codecs.droid ;; *) echo "Unknown vendor '$VENDOR'. Supported: lerobot_0_3_3 | lerobot | openpi | gr00t" >&2; exit 1 ;; esac @@ -146,9 +146,8 @@ case "$VENDOR" in "--input_path=$DATASET_DIR" \ "--output_path=$CKPT_DIR" \ "--exp_name=$EXP_NAME" \ - --num_train_steps=200 --save_steps=100 \ - --modality_config=ee_rot6d 2>&1) - SERVE_SUBCMD=(ee_rot6d --pipeline.source.checkpoints_dir="$CKPT_DIR$EXP_NAME/") + --num_train_steps=200 --save_steps=100 2>&1) + SERVE_SUBCMD=(droid --pipeline.source.checkpoints_dir="$CKPT_DIR$EXP_NAME/") ;; esac echo "$TRAIN_OUT" >> "$LOG" diff --git a/workflows/nebius/serve.sh b/workflows/nebius/serve.sh index 95abdcab8..28269868b 100644 --- a/workflows/nebius/serve.sh +++ b/workflows/nebius/serve.sh @@ -57,7 +57,7 @@ Examples: --pipeline.ee_frame=None # GR00T - bash workflows/nebius/serve.sh gr00t groot-server ee_rot6d_rel \ + bash workflows/nebius/serve.sh gr00t groot-server droid \ --pipeline.source.checkpoints_dir=s3:///checkpoints/groot// EOF exit 1 From 47b01d1ddf29bc12a7d0146ec7e5e8136546b082 Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Wed, 9 Sep 2026 21:21:15 +0300 Subject: [PATCH 02/14] Restore GR00T recording storage context --- positronic/vendors/gr00t/README.md | 4 ++++ positronic/vendors/gr00t/server.py | 3 ++- 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/positronic/vendors/gr00t/README.md b/positronic/vendors/gr00t/README.md index c2510ac7c..daf26f8f6 100644 --- a/positronic/vendors/gr00t/README.md +++ b/positronic/vendors/gr00t/README.md @@ -41,6 +41,10 @@ export IMAGE_TAG=local The GR00T environment is `/opt/gr00t-venv` (Python 3.12, upstream locked dependencies). Positronic has a separate environment at `/positronic/.venv`. Training and serving require a CUDA GPU. +The checkpoint also loads the gated `nvidia/Cosmos-Reason2-2B` backbone. The Hugging Face account +must have access to it, with its token available inside the container through `HF_TOKEN`, +`HF_TOKEN_PATH`, or the mounted Hugging Face cache's `token` file. + ## Convert and fine-tune From Positronic's `docker` directory: diff --git a/positronic/vendors/gr00t/server.py b/positronic/vendors/gr00t/server.py index 50b049c46..6d95e0c53 100644 --- a/positronic/vendors/gr00t/server.py +++ b/positronic/vendors/gr00t/server.py @@ -365,4 +365,5 @@ def pipeline(codec, source): if __name__ == '__main__': init_logging() - cfn.cli(COMMANDS) + with pos3.mirror(): + cfn.cli(COMMANDS) From 3a0fb78f3043e40197754b64f6382b9b2a59f5ed Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Wed, 9 Sep 2026 21:41:39 +0300 Subject: [PATCH 03/14] Keep uncharged inference within the simulation tick --- positronic/policy/harness.py | 17 +++++++++++++---- positronic/policy/layers.py | 8 ++++---- positronic/policy/tests/test_harness.py | 25 ++++++++++++++++++++++--- positronic/policy/tests/test_layers.py | 7 +++++++ 4 files changed, 46 insertions(+), 11 deletions(-) diff --git a/positronic/policy/harness.py b/positronic/policy/harness.py index 2a2ae2811..90223aa84 100644 --- a/positronic/policy/harness.py +++ b/positronic/policy/harness.py @@ -84,12 +84,22 @@ def _owned(obs: dict[str, Any]) -> dict[str, Any]: """ return {name: value.copy() if isinstance(value, np.ndarray) else value for name, value in obs.items()} - def __call__(self, obs: dict[str, Any]) -> list[dict[str, Any]] | None: + def __call__(self, obs: dict[str, Any], should_stop: pimm.SignalReceiver[bool]) -> list[dict[str, Any]] | None: now_ns = self._clock.now_ns() # A call that joins work already in flight keeps its anchor, so the trial pays for that work one time. if not self._rollout.rt.in_flight: self._t0_ns, self._wall_t0 = now_ns, time.monotonic() - return self._rollout.session(frozen_view(self._owned(obs)), now_ns) + owned = frozen_view(self._owned(obs)) + while True: + trajectory = self._rollout.session(owned, now_ns) + self.wait(should_stop) + if ( + trajectory is not None + or self._charges_wall_time + or not self._rollout.rt.owes_an_answer + or should_stop.value + ): + return trajectory def wait(self, should_stop: pimm.SignalReceiver[bool]) -> None: """Wait for the function in flight, for as long as the trial charges the loop for it.""" @@ -357,9 +367,8 @@ def _infer(self, inference: _EpisodeInference, clock: pimm.Clock, should_stop: p obs = self._build_obs(clock) except pimm.NoValueException: return # no function is in flight yet, so this skips no wait - if (trajectory := inference(obs)) is not None: + if (trajectory := inference(obs, should_stop)) is not None: self._reschedule(trajectory, clock) - inference.wait(should_stop) @staticmethod def _assert_anchored(trajectory: list[dict[str, Any]], now: float) -> None: diff --git a/positronic/policy/layers.py b/positronic/policy/layers.py index 530039e06..37c037754 100644 --- a/positronic/policy/layers.py +++ b/positronic/policy/layers.py @@ -78,10 +78,10 @@ class _Session(DelegatingSession): def __init__(self, inner: Session): super().__init__(inner) - self._trajectory_end: float | None = None + self._trajectory_end_ns: int | None = None def __call__(self, obs, time_ns): - if self._trajectory_end is not None and _obs_time(obs) < self._trajectory_end: + if self._trajectory_end_ns is not None and obs[keys.OBS_TIME_NS] < self._trajectory_end_ns: return None result = self._inner(obs, time_ns) if result is not None: @@ -93,11 +93,11 @@ def __call__(self, obs, time_ns): # Copy dicts so we don't mutate caller-owned data (sessions may reuse templates). anchor = time_ns / 1e9 result = [{**r, keys.ACTION_TIMESTAMP: anchor + r.get(keys.ACTION_TIMESTAMP, 0.0)} for r in result] - self._trajectory_end = result[-1][keys.ACTION_TIMESTAMP] if result else None + self._trajectory_end_ns = round(result[-1][keys.ACTION_TIMESTAMP] * 1e9) if result else None return result def cancel(self): - self._trajectory_end = None + self._trajectory_end_ns = None super().cancel() def make_session(self, inner: Session): diff --git a/positronic/policy/tests/test_harness.py b/positronic/policy/tests/test_harness.py index b05bed4c8..23c1d1f7b 100644 --- a/positronic/policy/tests/test_harness.py +++ b/positronic/policy/tests/test_harness.py @@ -710,12 +710,15 @@ def functions(self): rollout = Rollout(Task(instruction_source='t', timeout_sec=None), _HangingPolicy(), None) inference = _EpisodeInference(rollout, charges_wall_time=False, clock=world.clock) + stop = threading.Timer(0.02, world.request_stop) try: - inference({}) # starts the function, which never answers - world.request_stop() - inference.wait(world.should_stop_reader()) + stop.start() + assert inference({}, world.should_stop_reader()) is None + assert world.should_stop_reader().value finally: + stop.cancel() never_answers.set() + rollout.close() @pytest.mark.timeout(3.0) @@ -2124,6 +2127,22 @@ def test_an_uncharged_call_pauses_the_world(world): assert played[0][0] < 0.05, f'the world paid for the function: first command at {played[0][0]}s' +@pytest.mark.timeout(20.0) +@pytest.mark.parametrize('wall_sec', [0.0, 0.01]) +def test_uncharged_chunks_have_no_extra_control_tick(world, wall_sec): + chunk = [*slow_chunk(0.1, 4), {keys.ACTION_TIMESTAMP: 0.1}] + played = _run_episode( + world, + RemoteStubPolicy(wall_sec=wall_sec, chunk=chunk), + ChunkedSchedule(), + charge_inference_time=False, + run_sec=0.4, + ) + + assert len(played) >= 12 + np.testing.assert_allclose(np.diff([t for t, _ in played])[3::4], 0.025, atol=1e-7) + + @pytest.mark.timeout(20.0) def test_a_charged_call_costs_its_own_wall_duration(world): """``charge_inference_time=True`` charges the world what the model really took, so a slow server is scored diff --git a/positronic/policy/tests/test_layers.py b/positronic/policy/tests/test_layers.py index 28a5d1007..164ceb1b7 100644 --- a/positronic/policy/tests/test_layers.py +++ b/positronic/policy/tests/test_layers.py @@ -188,6 +188,13 @@ def test_single_action_refires_immediately_after(self): result = session(_obs(1.01), int(1.01e9)) assert result is not None + def test_chunk_expiry_rounds_to_the_clock_nanosecond(self): + session = ChunkedSchedule().make_session(_ConstSession([{keys.ACTION_TIMESTAMP: 0.1}])) + session(_obs(0.2), 200_000_000) + + assert session({keys.OBS_TIME_NS: 299_999_999}, 299_999_999) is None + assert session({keys.OBS_TIME_NS: 300_000_000}, 300_000_000) is not None + def test_expiry_is_judged_at_the_observation_instant(self): """Whether the trajectory has run out is a question about the observation, not about the call's time.""" inner = _ConstPolicy([{'v': 1, keys.ACTION_TIMESTAMP: 0.0}, {'v': 2, keys.ACTION_TIMESTAMP: 0.5}]) From ebd57ec68554ed264042cc86951048ed15e5a3e2 Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Thu, 10 Sep 2026 12:32:40 +0300 Subject: [PATCH 04/14] Polish GR00T runtime and cloud workflows --- docker/Dockerfile.groot | 6 +- docker/Makefile | 12 +- docker/README.md | 13 +- docs/codecs.md | 6 +- docs/inference.md | 4 +- docs/training-workflow.md | 6 +- positronic/policy/layers.py | 7 +- positronic/policy/tests/test_harness.py | 36 +++--- positronic/vendors/gr00t/README.md | 28 +++-- positronic/vendors/gr00t/server.py | 115 ++++++++---------- positronic/vendors/gr00t/tests/test_server.py | 22 ++++ positronic/vendors/gr00t/tests/test_train.py | 2 - positronic/vendors/gr00t/train.py | 70 ++++++----- workflows/nebius/README.md | 5 + workflows/nebius/common.sh | 5 + workflows/nebius/serve.sh | 9 +- workflows/nebius/train.sh | 9 +- 17 files changed, 188 insertions(+), 167 deletions(-) diff --git a/docker/Dockerfile.groot b/docker/Dockerfile.groot index ee9a6a315..5a2bdf866 100644 --- a/docker/Dockerfile.groot +++ b/docker/Dockerfile.groot @@ -1,22 +1,18 @@ -ARG BASE_IMAGE=positro/gr00t-base:240627d +ARG BASE_IMAGE FROM ${BASE_IMAGE} COPY --from=ghcr.io/astral-sh/uv:latest /uv /uvx /usr/local/bin/ -# Set working directory WORKDIR /positronic RUN apt-get update && apt-get install -y --no-install-recommends libturbojpeg && rm -rf /var/lib/apt/lists/* -# Copy the positronic repository COPY . /positronic -# Set Python path to include source root and gr00t ENV PYTHONPATH=/positronic:/gr00t # Keep Positronic's dependencies separate from the checkpoint runtime. ENV UV_PROJECT_ENVIRONMENT=/positronic/.venv VIRTUAL_ENV=/positronic/.venv RUN --mount=type=cache,target=/root/.cache/uv uv sync --locked --python 3.12 --no-dev -# Default command CMD ["bash"] diff --git a/docker/Makefile b/docker/Makefile index b4c082d0f..82a27c84d 100644 --- a/docker/Makefile +++ b/docker/Makefile @@ -1,5 +1,7 @@ .PHONY: all build tag push clean prune help build-training tag-training push-training build-openpi tag-openpi push-openpi build-groot tag-groot push-groot build-dreamzero-base push-dreamzero-base build-dreamzero tag-dreamzero push-dreamzero build-robolab tag-robolab push-robolab nebius-login push-robolab-cr +.DEFAULT_GOAL := help + # Image configuration IMAGE_NAME_TRAINING := positro/positronic IMAGE_NAME_OPENPI := positro/openpi @@ -41,7 +43,15 @@ TAG_DREAMZERO_BASE_LATEST := $(IMAGE_NAME_DREAMZERO_BASE):latest # OpenPI base image (pulled from registry) TAG_OPENPI_BASE := $(IMAGE_NAME_OPENPI_BASE):latest -GROOT_BASE_IMAGE ?= positro/gr00t-base:240627d +GROOT_REF := 506ccc87314fb69ad7ac8d8960b002ce3a378395 +ifndef GROOT_BASE_IMAGE +GROOT_BASE_IMAGE := positro/gr00t-base:$(GROOT_REF) +build-groot: build-groot-base +endif + +.PHONY: build-groot-base +build-groot-base: + docker build --platform linux/amd64 -f docker/Dockerfile -t $(GROOT_BASE_IMAGE) https://github.com/Positronic-Robotics/gr00t.git\#$(GROOT_REF) help: @echo "Positronic Docker Build System - Makefile" diff --git a/docker/README.md b/docker/README.md index a43e0ab9f..02d290c72 100644 --- a/docker/README.md +++ b/docker/README.md @@ -42,20 +42,21 @@ docker/build.sh ## GR00T N1.7 containers -`make build-groot` uses `GROOT_BASE_IMAGE` from the GR00T fork. The fork image contains +`make build-groot` builds the GR00T fork revision pinned in `Makefile`. The fork image contains CUDA 12.8 and the upstream Python 3.12 environment at `/opt/gr00t-venv`. Positronic installs its own locked environment at `/positronic/.venv`. Both training and serving launch GR00T in its separate environment. -Build the base in the [fork](https://github.com/Positronic-Robotics/gr00t), then build the adapter: +Build both images from Positronic: ```bash -# In the GR00T fork: -make -C docker build -# In Positronic: -make -C docker build-groot GROOT_BASE_IMAGE=positro/gr00t-base:local +make -C docker build-groot IMAGE_TAG=local docker compose -f docker/docker-compose.yml run --rm --service-ports groot-server droid ``` +Pass `GROOT_BASE_IMAGE=` to use an existing base and skip its build. +For local fork development, run `make -C docker build` in the fork, +then `make -C docker build-groot GROOT_BASE_IMAGE=positro/gr00t-base:local` in Positronic. + Do not mount host uv interpreter directories over the image's interpreter directories. See [GR00T](../positronic/vendors/gr00t/README.md) for conversion, fine-tuning and inference. diff --git a/docs/codecs.md b/docs/codecs.md index 3aef6fd98..aa790a89c 100644 --- a/docs/codecs.md +++ b/docs/codecs.md @@ -57,7 +57,7 @@ It is declared in two places, for the two things it does: - **Training** — `compose(ee_frame=DROID_EE_FRAME)` re-expresses the dataset in that frame, which is what makes the resulting checkpoint speak it. It defaults to unset, which trains in `default`. - **Serving** — the OpenPI pipeline's `ee_frame=` puts the codec left of the `remote` marker, so the rig converts and the server stays frame-agnostic. It has no default: every deployment states its frame — `None` for a checkpoint trained in `default`, or one that speaks joints, which are unambiguous. Nothing checks a stated frame against how the checkpoint was trained, so it is set beside the checkpoint path it belongs to. -Both take the transform itself — `models.DROID_EE_FRAME` is the one we ship — so a checkpoint declares its own frame and no robot model is consulted to serve it. The other vendor servers take no `ee_frame`: every checkpoint they serve was trained in the rig's `default`, so none has a transform to declare. +Both take the transform itself — `models.DROID_EE_FRAME` is the one we ship — so a checkpoint declares its own frame and no robot model is consulted to serve it. GR00T's DROID codec uses `DROID_EE_FRAME` for both dataset conversion and serving. Its pipeline exposes this through `codec.ee_frame`. A `CartesianDelta` is the one command this cannot convert on its own: a delta has no anchor pose, so it carries `frame` and the driver composes it where the measured pose lives. @@ -70,7 +70,9 @@ Two wrappers in [`positronic/cfg/codecs.py`](../positronic/cfg/codecs.py) apply | Wrapper | Expands to | Used by | |---------|-----------|---------| | `droid_execution(action)` | `SetControlMode(DROID_IMPEDANCE) \| action` ([the DROID gains](../positronic/cfg/hardware/roboarm/__init__.py)) | the `droid` pipelines of OpenPI, DreamZero and MolmoAct2, and OpenPI's `droid_jointpos` | -| `phail_v1_execution(action)` | `SetControlMode(PositionControl()) \| action` | the `phail_v1` pipelines of LeRobot, GR00T, OpenPI and DreamZero | +| `phail_v1_execution(action)` | `SetControlMode(PositionControl()) \| action` | the `phail_v1` pipelines of LeRobot, OpenPI and DreamZero | + +GR00T's DROID codec sets `DROID_IMPEDANCE` directly on its joint-position commands. ## Writing custom codecs diff --git a/docs/inference.md b/docs/inference.md index c8b1ea714..dd7c6b4dc 100644 --- a/docs/inference.md +++ b/docs/inference.md @@ -20,8 +20,8 @@ cd docker && docker compose run --rm --service-ports lerobot-0_3_3-server ee \ --pipeline.source.checkpoints_dir=~/checkpoints/lerobot/experiment_v1/ # GR00T -cd docker && docker compose run --rm --service-ports groot-server droid \ - --pipeline.source.checkpoints_dir=~/checkpoints/groot/experiment_v1/ +cd docker && docker compose run --rm --service-ports -v "$PWD/groot-data:/data" groot-server droid \ + --pipeline.source.checkpoints_dir=/data/checkpoints/experiment_v1/ # OpenPI (--pipeline.ee_frame states the EE frame the checkpoint speaks; None means the rig's `default`) cd docker && docker compose run --rm --service-ports openpi-server ee \ diff --git a/docs/training-workflow.md b/docs/training-workflow.md index e9bbede82..545ab09bb 100644 --- a/docs/training-workflow.md +++ b/docs/training-workflow.md @@ -124,9 +124,9 @@ cd docker && docker compose run --rm lerobot-train full_finetune \ ### GR00T Training ```bash -cd docker && docker compose run --rm groot-train \ - --input_path=~/datasets/groot/stack_cubes \ - --output_path=~/checkpoints/groot \ +cd docker && docker compose run --rm -v "$PWD/groot-data:/data" groot-train \ + --input_path=/data/datasets/stack_cubes \ + --output_path=/data/checkpoints \ --exp_name=experiment_v1 ``` diff --git a/positronic/policy/layers.py b/positronic/policy/layers.py index 37c037754..3a6c3be23 100644 --- a/positronic/policy/layers.py +++ b/positronic/policy/layers.py @@ -14,11 +14,6 @@ from positronic.policy.base import DelegatingSession, Layer, Session -def _obs_time(obs) -> float: - """Observation timestamp in seconds, from the harness's nanosecond stamp.""" - return obs[keys.OBS_TIME_NS] / 1e9 - - # TODO(#638): the arm is found by name because the harness serializes before the stack sees anything. Once # domain types reach the border, this reads the status off the value. def _is_robot_status(name: str) -> bool: @@ -179,7 +174,7 @@ def __init__(self, inner: Session, keys: tuple[str, ...], offsets_sec: tuple[flo self._buffer = _StackBuffer(offsets_sec, pad_start=pad_start) def __call__(self, obs, time_ns): - now = _obs_time(obs) + now = obs[keys.OBS_TIME_NS] / 1e9 self._buffer.append(now, {k: obs[k] for k in self._keys}) return self._inner({**obs, **self._buffer.sample(now)}, time_ns) diff --git a/positronic/policy/tests/test_harness.py b/positronic/policy/tests/test_harness.py index 23c1d1f7b..3aab2a3ee 100644 --- a/positronic/policy/tests/test_harness.py +++ b/positronic/policy/tests/test_harness.py @@ -339,14 +339,8 @@ def _last_grip(p): return msg.data -def _emitted_commands(recorder): - """Every robot command a recorder saw, in emission order.""" - return [cmd for _ts, cmd in recorder.emitted] - - -def _emitted_grips(recorder): - """Every grip target a recorder saw, in emission order.""" - return [grip for _ts, grip in recorder.emitted] +def _emitted_values(recorder): + return [value for _ts, value in recorder.emitted] @pytest.mark.timeout(3.0) @@ -399,14 +393,14 @@ def test_harness_emits_cartesian_move(world): keys.DESCRIPTOR, } - cmds = _emitted_commands(cmd_recorder) + cmds = _emitted_values(cmd_recorder) assert cmds, 'no robot command emitted' cmd = cmds[-1] assert isinstance(cmd, roboarm.command.CartesianPosition) np.testing.assert_allclose(cmd.pose.translation, pose.translation) np.testing.assert_allclose(cmd.pose.rotation.as_quat, pose.rotation.as_quat) - grips = _emitted_grips(grip_recorder) + grips = _emitted_values(grip_recorder) assert grips and grips[-1] == pytest.approx(0.33) @@ -538,7 +532,7 @@ def test_harness_waits_for_complete_inputs(world): def assert_no_inference(): assert policy.last_obs is None - assert not _emitted_commands(cmd_recorder) + assert not _emitted_values(cmd_recorder) driver = ManualDriver([ (partial(perform_task, Task(instruction_source='dummy-task', timeout_sec=None)), 0.01), @@ -554,13 +548,13 @@ def assert_no_inference(): assert policy.last_obs is not None - cmds = _emitted_commands(cmd_recorder) + cmds = _emitted_values(cmd_recorder) assert cmds, 'no robot command emitted' cmd = cmds[-1] assert isinstance(cmd, roboarm.command.CartesianPosition) np.testing.assert_allclose(cmd.pose.translation, pose.translation) - grips = _emitted_grips(grip_recorder) + grips = _emitted_values(grip_recorder) assert grips and grips[-1] == pytest.approx(0.33) @@ -1370,8 +1364,8 @@ def test_timeout_during_inference_drops_the_chunk(world): assert len(stops) == 1 assert stops[0].static_data[eval_keys.TERMINATED] is False # A trial that times out mid-call plays nothing: the chunk it was waiting on is dropped. - assert not _emitted_commands(cmd_recorder) - assert not _emitted_grips(grip_recorder) + assert not _emitted_values(cmd_recorder) + assert not _emitted_values(grip_recorder) @pytest.mark.timeout(3.0) @@ -1635,8 +1629,8 @@ def new_session(self, context=None, rt=None): scheduler = world.start([harness, ManualDriver(script)]) drive_scheduler(scheduler, steps=200) - assert not _emitted_commands(cmd_recorder) # an empty trajectory schedules nothing - assert not _emitted_grips(grip_recorder) + assert not _emitted_values(cmd_recorder) # an empty trajectory schedules nothing + assert not _emitted_values(grip_recorder) @pytest.mark.timeout(3.0) @@ -2503,7 +2497,7 @@ def new_session(self, context=None, rt=None): ]) drive_scheduler(world.start([harness, driver]), steps=1000) - grips = _emitted_grips(grip_recorder) + grips = _emitted_values(grip_recorder) assert set(grips) == {0.5}, f'the second chunk kept the gripper playing: {grips}' @@ -2523,8 +2517,8 @@ def test_manual_commands_are_emitted_as_plain_values(world): driver = ManualDriver([(partial(manual_em.emit, {keys.ROBOT_COMMAND: manual}), 0.01), (None, 0.02)]) drive_scheduler(world.start([harness, driver]), steps=50) - assert _emitted_commands(cmd_recorder) == [manual] - assert not _emitted_grips(grip_recorder) + assert _emitted_values(cmd_recorder) == [manual] + assert not _emitted_values(grip_recorder) @pytest.mark.timeout(20.0) @@ -2554,4 +2548,4 @@ def test_finishing_discards_a_call_that_is_still_in_flight(world): ]) drive_scheduler(world.start([harness, driver, _Pacer()]), steps=2000) - assert not _emitted_commands(cmd_recorder) + assert not _emitted_values(cmd_recorder) diff --git a/positronic/vendors/gr00t/README.md b/positronic/vendors/gr00t/README.md index daf26f8f6..2509cd032 100644 --- a/positronic/vendors/gr00t/README.md +++ b/positronic/vendors/gr00t/README.md @@ -27,17 +27,17 @@ N1.6 checkpoints require an N1.6 image; their custom action schemas are incompat ## Docker -Build the fork's base image, then the Positronic image: +Build the pinned fork and the Positronic image: ```bash -# In the GR00T fork -make -C docker build -# In Positronic -make -C docker build-groot GROOT_BASE_IMAGE=positro/gr00t-base:local +make -C docker build-groot cd docker export IMAGE_TAG=local ``` +`build-groot` builds its base from the fork revision pinned in `docker/Makefile`. +To use an existing base image, pass `GROOT_BASE_IMAGE=`. + The GR00T environment is `/opt/gr00t-venv` (Python 3.12, upstream locked dependencies). Positronic has a separate environment at `/positronic/.venv`. Training and serving require a CUDA GPU. @@ -50,13 +50,14 @@ must have access to it, with its token available inside the container through `H From Positronic's `docker` directory: ```bash -docker compose run --rm --pull never lerobot-0_3_3-convert convert \ +mkdir -p "$PWD/groot-data" +docker compose run --rm --pull never -v "$PWD/groot-data:/data" lerobot-0_3_3-convert convert \ --dataset.codec=@positronic.vendors.gr00t.codecs.droid \ - --output_dir=~/datasets/groot/my_task + --output_dir=/data/datasets/my_task -docker compose run --rm groot-train \ - --input_path=~/datasets/groot/my_task \ - --output_path=~/checkpoints/groot \ +docker compose run --rm --pull never -v "$PWD/groot-data:/data" groot-train \ + --input_path=/data/datasets/my_task \ + --output_path=/data/checkpoints \ --exp_name=my_task \ --num_train_steps=10000 ``` @@ -82,8 +83,8 @@ docker compose run --rm --service-ports groot-server droid Fine-tuned checkpoint: ```bash -docker compose run --rm --service-ports groot-server droid \ - --pipeline.source.checkpoints_dir=~/checkpoints/groot/my_task +docker compose run --rm --service-ports --pull never -v "$PWD/groot-data:/data" groot-server droid \ + --pipeline.source.checkpoints_dir=/data/checkpoints/my_task ``` Select `droid_three_cameras` for a checkpoint trained on three views. @@ -95,7 +96,8 @@ A Hugging Face source uses `--pipeline.source.checkpoints_dir=hf://owner/model`. The GR00T source is included in the image at `/gr00t`. From the image's `/positronic` directory: ```bash -GR00T_REFERENCE_ROOT=/gr00t uv run --no-sync --python 3.12 pytest \ +uv sync --locked --python 3.12 +GR00T_REFERENCE_ROOT=/gr00t uv run --no-sync --python 3.12 python -m pytest \ -o addopts= positronic/vendors/gr00t/tests ``` diff --git a/positronic/vendors/gr00t/server.py b/positronic/vendors/gr00t/server.py index 6d95e0c53..c0bdf2028 100644 --- a/positronic/vendors/gr00t/server.py +++ b/positronic/vendors/gr00t/server.py @@ -62,37 +62,34 @@ def __init__(self, host: str = 'localhost', port: int = 5555, timeout_ms: int = self.host = host self.port = port self.timeout_ms = timeout_ms - self._init_socket() + self.socket = self._make_socket() - def _init_socket(self): - if hasattr(self, 'socket'): - self.socket.close(linger=0) - self.socket = self.context.socket(zmq.REQ) - self.socket.setsockopt(zmq.RCVTIMEO, self.timeout_ms) - self.socket.setsockopt(zmq.SNDTIMEO, self.timeout_ms) - self.socket.connect(f'tcp://{self.host}:{self.port}') + def _make_socket(self): + socket = self.context.socket(zmq.REQ) + socket.setsockopt(zmq.RCVTIMEO, self.timeout_ms) + socket.setsockopt(zmq.SNDTIMEO, self.timeout_ms) + socket.connect(f'tcp://{self.host}:{self.port}') + return socket def ping(self) -> bool: try: - self.call_endpoint('ping', requires_input=False) + self.call_endpoint('ping') return True except (zmq.error.ZMQError, RuntimeError): - self._init_socket() return False - def call_endpoint(self, endpoint: str, data: dict | None = None, requires_input: bool = True) -> Any: + def call_endpoint(self, endpoint: str, data: dict | None = None) -> Any: request: dict = {'endpoint': endpoint} - if requires_input: + if data is not None: request['data'] = data try: self.socket.send(MsgSerializer.to_bytes(request)) message = self.socket.recv() - except zmq.error.Again as err: - self._init_socket() - raise RuntimeError( - f'Timeout after {self.timeout_ms}ms calling endpoint "{endpoint}" at {self.host}:{self.port}' - ) from err + except zmq.error.ZMQError as err: + self.socket.close(linger=0) + self.socket = self._make_socket() + raise RuntimeError(f'GR00T endpoint {endpoint} failed at {self.host}:{self.port}: {err}') from err if message == b'ERROR': raise RuntimeError('Server error. Make sure the correct policy server is running.') @@ -114,15 +111,10 @@ def close(self): self.context.term() -########################################################################################### -# Subprocess manager for gr00t ZMQ server -########################################################################################### - - class Gr00tSubprocess: """Manages the gr00t ZMQ server subprocess.""" - def __init__(self, checkpoint_dir: str, groot_venv_path: str, zmq_port: int = 5555, ready_timeout: float = 120.0): + def __init__(self, checkpoint_dir: str, groot_venv_path: Path, zmq_port: int = 5555, ready_timeout: float = 120.0): self.checkpoint_dir = checkpoint_dir self.groot_venv_path = groot_venv_path self.zmq_port = zmq_port @@ -132,7 +124,7 @@ def __init__(self, checkpoint_dir: str, groot_venv_path: str, zmq_port: int = 55 def start(self, on_progress: Callable[[str], None] | None = None): groot_root = Path(__file__).parents[4] / 'gr00t' - python_bin = str(Path(self.groot_venv_path) / 'bin' / 'python') + python_bin = str(self.groot_venv_path / 'bin' / 'python') command = [python_bin, 'gr00t/eval/run_gr00t_server.py'] command.extend(['--model_path', str(self.checkpoint_dir)]) @@ -178,14 +170,10 @@ def stop(self): self.process.wait(timeout=10) except subprocess.TimeoutExpired: self.process.kill() + self.process.wait() self.process = None -########################################################################################### -# Policy and model source -########################################################################################### - - class _Gr00tSession(Session): def __init__(self, client: PolicyClient, meta: dict[str, Any]): self._client = client @@ -219,32 +207,6 @@ def close(self): self._groot.stop() -def _step_id(raw: str) -> str: - """The public id for a ``checkpoint-`` directory: its step number, free of any zero-padding.""" - return str(int(raw)) if raw.isdigit() else raw - - -def _warm_observation(modalities: dict) -> dict[str, Any]: - """Build a valid current-frame DROID observation from the loaded checkpoint's modalities.""" - for name in (gr00t.VIDEO, gr00t.STATE): - if modalities[name]['delta_indices'] != [0]: - raise ValueError(f'DROID adapter requires current-frame {name}, got {modalities[name]}') - width, height = gr00t.IMAGE_SIZE - state = { - name: np.zeros((1, 1, gr00t.STATE_DIMS[name]), dtype=np.float32) - for name in modalities[gr00t.STATE]['modality_keys'] - } - state[gr00t.EE_POSE][..., 3:] = [1, 0, 0, 0, 1, 0] - return { - gr00t.VIDEO: { - name: np.zeros((1, 1, height, width, 3), dtype=np.uint8) - for name in modalities[gr00t.VIDEO]['modality_keys'] - }, - gr00t.STATE: state, - gr00t.LANGUAGE: {gr00t.TASK: [['pick up the object']]}, - } - - class Gr00tSource(ModelSource): """A Hugging Face model (``hf://owner/model``) or a directory of fine-tuned checkpoints. @@ -264,7 +226,7 @@ def __init__( if self._is_hub_model and checkpoint is not None: raise ValueError('checkpoint step selection applies only to fine-tuned checkpoint directories') self.checkpoint = checkpoint - self.groot_venv_path = groot_venv_path + self.groot_venv_path = Path(groot_venv_path).expanduser() self.zmq_port = zmq_port self.ready_timeout = ready_timeout @@ -282,10 +244,15 @@ def _raw_for(self, model_id: str) -> str: def _is_hub_model(self) -> bool: return self.checkpoints_dir.startswith('hf://') + @staticmethod + def _step_id(raw: str) -> str: + """The public id for a ``checkpoint-`` directory: its step number, free of any zero-padding.""" + return str(int(raw)) if raw.isdigit() else raw + def get_models(self) -> list[str]: if self._is_hub_model: return [self.checkpoints_dir.removeprefix('hf://')] - return [_step_id(r) for r in self._raw_ids()] + return [self._step_id(r) for r in self._raw_ids()] def resolve(self, model_id: str | None) -> str: """Explicit id > the configured ``checkpoint`` > latest, always as the id ``get_models`` advertises. @@ -301,8 +268,29 @@ def resolve(self, model_id: str | None) -> str: if model_id is None and self.checkpoint is not None: model_id = str(self.checkpoint).strip('/') if model_id is None: - return _step_id(self._raw_ids()[-1]) - return _step_id(self._raw_for(model_id)) + return self._step_id(self._raw_ids()[-1]) + return self._step_id(self._raw_for(model_id)) + + @staticmethod + def _warm_observation(modalities: dict) -> dict[str, Any]: + """Build a valid current-frame DROID observation from the loaded checkpoint's modalities.""" + for name in (gr00t.VIDEO, gr00t.STATE): + if modalities[name]['delta_indices'] != [0]: + raise ValueError(f'DROID adapter requires current-frame {name}, got {modalities[name]}') + width, height = gr00t.IMAGE_SIZE + state = { + name: np.zeros((1, 1, gr00t.STATE_DIMS[name]), dtype=np.float32) + for name in modalities[gr00t.STATE]['modality_keys'] + } + state[gr00t.EE_POSE][..., 3:] = [1, 0, 0, 0, 1, 0] + return { + gr00t.VIDEO: { + name: np.zeros((1, 1, height, width, 3), dtype=np.uint8) + for name in modalities[gr00t.VIDEO]['modality_keys'] + }, + gr00t.STATE: state, + gr00t.LANGUAGE: {gr00t.TASK: [['pick up the object']]}, + } def load(self, model_id: str, on_progress: Callable[[str], None] | None = None) -> Policy: if self._is_hub_model: @@ -325,8 +313,8 @@ def load(self, model_id: str, on_progress: Callable[[str], None] | None = None) groot.start(on_progress) policy = Gr00tPolicy(groot, str(checkpoint_dir)) # The subprocess initializes CUDA on its first forward, which outlasts a rig's inference timeout. - modalities = groot.client.call_endpoint('get_modality_config', requires_input=False) - warmup(policy, _warm_observation(modalities), on_progress) + modalities = groot.client.call_endpoint('get_modality_config') + warmup(policy, self._warm_observation(modalities), on_progress) except Exception: groot.stop() raise @@ -340,11 +328,6 @@ def meta(self, model_id: str) -> dict[str, Any]: } -########################################################################################### -# Serving configs -########################################################################################### - - gr00t_source = cfn.Config(Gr00tSource) diff --git a/positronic/vendors/gr00t/tests/test_server.py b/positronic/vendors/gr00t/tests/test_server.py index c0484308c..72ccff287 100644 --- a/positronic/vendors/gr00t/tests/test_server.py +++ b/positronic/vendors/gr00t/tests/test_server.py @@ -1,7 +1,10 @@ +from unittest.mock import Mock + import msgpack import msgpack_numpy import numpy as np import pytest +import zmq from positronic.vendors import gr00t from positronic.vendors.gr00t import server as gr00t_server @@ -49,3 +52,22 @@ def test_published_checkpoint_is_served_without_a_local_checkpoint_scan(monkeypa source = gr00t_server.Gr00tSource() assert source.get_models() == [gr00t.BASE_MODEL] assert source.resolve(None) == gr00t.BASE_MODEL + + +@pytest.mark.parametrize('failure', [zmq.Again(), zmq.ZMQError(zmq.EFSM)]) +def test_client_can_ping_after_a_transport_failure(monkeypatch, failure): + failed = Mock() + failed.send.side_effect = failure + recovered = Mock() + recovered.recv.return_value = gr00t_server.MsgSerializer.to_bytes('pong') + context = Mock() + context.socket.side_effect = [failed, recovered] + monkeypatch.setattr(gr00t_server.zmq, 'Context', lambda: context) + client = gr00t_server.PolicyClient() + try: + assert not client.ping() + assert client.ping() + request = gr00t_server.MsgSerializer.from_bytes(recovered.send.call_args.args[0]) + assert request == {'endpoint': 'ping'} + finally: + client.close() diff --git a/positronic/vendors/gr00t/tests/test_train.py b/positronic/vendors/gr00t/tests/test_train.py index c010120ee..c792fa458 100644 --- a/positronic/vendors/gr00t/tests/test_train.py +++ b/positronic/vendors/gr00t/tests/test_train.py @@ -1,5 +1,4 @@ import json -from contextlib import nullcontext from unittest.mock import Mock import pytest @@ -18,7 +17,6 @@ def test_finetuning_forwards_dataset_cameras_and_resume_to_gr00t(tmp_path, monke output.mkdir() sync = Mock(return_value=output) run = Mock() - monkeypatch.setattr(train.pos3, 'mirror', nullcontext) monkeypatch.setattr(train.pos3, 'download', lambda _: dataset) monkeypatch.setattr(train.pos3, 'sync', sync) monkeypatch.setattr(train.utils, 'save_run_metadata', Mock()) diff --git a/positronic/vendors/gr00t/train.py b/positronic/vendors/gr00t/train.py index 27d4c5a97..f1129a91c 100644 --- a/positronic/vendors/gr00t/train.py +++ b/positronic/vendors/gr00t/train.py @@ -10,9 +10,9 @@ from positronic.vendors import gr00t -def cleanup_old_optimizers(output_dir: str, keep_last_n: int = 2): +def cleanup_old_optimizers(output_dir: Path, keep_last_n: int = 2): """Delete optimizer.pt from all but the last N checkpoints to save space.""" - checkpoints = sorted(Path(output_dir).glob('checkpoint-*'), key=lambda p: int(p.name.split('-')[1])) + checkpoints = sorted(output_dir.glob('checkpoint-*'), key=lambda p: int(p.name.split('-')[1])) for ckpt in checkpoints[:-keep_last_n] if keep_last_n > 0 else checkpoints: opt_file = ckpt / 'optimizer.pt' if opt_file.exists(): @@ -21,6 +21,7 @@ def cleanup_old_optimizers(output_dir: str, keep_last_n: int = 2): @cfn.config(num_train_steps=None, groot_venv_path=gr00t.VENV, base_model=gr00t.BASE_MODEL, batch_size=64) +@pos3.mirror() def main( input_path: str, output_path: str, @@ -39,43 +40,40 @@ def main( groot_root = Path(__file__).parents[4] / 'gr00t' python_bin = str(Path(groot_venv_path).expanduser() / 'bin' / 'python') - with pos3.mirror(): - dataset_local_path = pos3.download(input_path) - with (Path(dataset_local_path) / 'meta' / 'modality.json').open() as f: - video_keys = list(json.load(f)[gr00t.VIDEO]) - output_path = output_path.rstrip('/') - # When resuming, don't delete existing checkpoint files - output_dir = pos3.sync(output_path + '/' + exp_name, delete_remote=not resume) - prefix = 'resume_metadata' if resume else 'run_metadata' - utils.save_run_metadata(output_dir, patterns=['*.py', '*.toml'], prefix=prefix) + dataset_local_path = pos3.download(input_path) + with (Path(dataset_local_path) / 'meta' / 'modality.json').open() as f: + video_keys = list(json.load(f)[gr00t.VIDEO]) + output_path = output_path.rstrip('/') + output_dir = pos3.sync(output_path + '/' + exp_name, delete_remote=not resume) + prefix = 'resume_metadata' if resume else 'run_metadata' + utils.save_run_metadata(output_dir, patterns=['*.py', '*.toml'], prefix=prefix) - command = [python_bin, 'gr00t/experiment/launch_finetune.py'] - command.extend(['--base-model-path', base_model]) - command.extend(['--dataset_path', str(dataset_local_path)]) - command.extend(['--video-keys', *video_keys]) - command.extend(['--embodiment-tag', gr00t.EMBODIMENT]) - command.extend(['--global-batch-size', str(batch_size)]) - if resume: - command.append('--resume-from-checkpoint') - command.extend(['--output_dir', str(output_dir)]) - command.extend(['--num_gpus', '1']) - command.extend(['--save_total_limit', '9999']) # Keep all checkpoints - if num_train_steps is not None: - command.extend(['--max_steps', str(num_train_steps)]) - if learning_rate is not None: - command.extend(['--learning_rate', str(learning_rate)]) - if save_steps is not None: - command.extend(['--save_steps', str(save_steps)]) - if num_workers is not None: - command.extend(['--dataloader_num_workers', str(num_workers)]) - command.append('--use-wandb') + command = [python_bin, 'gr00t/experiment/launch_finetune.py'] + command.extend(['--base-model-path', base_model]) + command.extend(['--dataset_path', str(dataset_local_path)]) + command.extend(['--video-keys', *video_keys]) + command.extend(['--embodiment-tag', gr00t.EMBODIMENT]) + command.extend(['--global-batch-size', str(batch_size)]) + if resume: + command.append('--resume-from-checkpoint') + command.extend(['--output_dir', str(output_dir)]) + command.extend(['--num_gpus', '1']) + command.extend(['--save_total_limit', '9999']) # Keep all checkpoints + if num_train_steps is not None: + command.extend(['--max_steps', str(num_train_steps)]) + if learning_rate is not None: + command.extend(['--learning_rate', str(learning_rate)]) + if save_steps is not None: + command.extend(['--save_steps', str(save_steps)]) + if num_workers is not None: + command.extend(['--dataloader_num_workers', str(num_workers)]) + command.append('--use-wandb') - env = os.environ.copy() - print(f'Running command: {command}') - subprocess.run(command, check=True, cwd=str(groot_root), env=env) + env = os.environ.copy() + print(f'Running command: {command}') + subprocess.run(command, check=True, cwd=str(groot_root), env=env) - # Clean up optimizer state from old checkpoints to save space - cleanup_old_optimizers(output_dir, keep_last_n=keep_optimizers_for_last_n) + cleanup_old_optimizers(Path(output_dir), keep_last_n=keep_optimizers_for_last_n) if __name__ == '__main__': diff --git a/workflows/nebius/README.md b/workflows/nebius/README.md index 8160bcd1e..5f53c5ff8 100644 --- a/workflows/nebius/README.md +++ b/workflows/nebius/README.md @@ -12,6 +12,11 @@ running inference from your robot or simulator against the served policy — wor reached at a managed `https://` URL and gated on a bearer token: see [Authenticated inference](#authenticated-inference). +GR00T training and serving use the image's Python 3.12 environment without resynchronizing dependencies. +They inject `HF_TOKEN` from the MysteryBox secret `huggingface-read-token`. +Override its name with `NEBIUS_HF_TOKEN_SECRET`; the payload key must be `HF_TOKEN`. +The token's account must have access to the gated `nvidia/Cosmos-Reason2-2B` backbone. + ## Prerequisites - Nebius CLI v0.12.209 or newer, authenticated to your project diff --git a/workflows/nebius/common.sh b/workflows/nebius/common.sh index dfc68ddb4..4d7218137 100644 --- a/workflows/nebius/common.sh +++ b/workflows/nebius/common.sh @@ -29,6 +29,11 @@ IMAGE_TAG="${NEBIUS_IMAGE_TAG:-latest}" AUTH_TOKEN_SECRET="${NEBIUS_AUTH_TOKEN_SECRET:-positronic-serverless-inference-token}" AUTH_TOKEN_KEY=AUTH_TOKEN +# The Hugging Face account must have access to the checkpoint's gated backbone. +HF_TOKEN_SECRET="${NEBIUS_HF_TOKEN_SECRET:-huggingface-read-token}" +HF_ENV_FLAGS=(--env-secret "HF_TOKEN=${HF_TOKEN_SECRET}") +GR00T_UV_ARGS="run --no-sync --python 3.12" + # S3 credentials and endpoint for pos3, as `nebius ai job|endpoint create` flags. Expand into a # create call with "${S3_ENV_FLAGS[@]}". S3_ENV_FLAGS=( diff --git a/workflows/nebius/serve.sh b/workflows/nebius/serve.sh index 28269868b..15a863898 100644 --- a/workflows/nebius/serve.sh +++ b/workflows/nebius/serve.sh @@ -67,12 +67,15 @@ VENDOR="$1" NAME="$2" shift 2 +UV_ARGS="run --python 3.13" +HF_FLAGS=() case "$VENDOR" in lerobot_0_3_3) IMAGE="positro/positronic:${IMAGE_TAG}"; EXTRA="--extra lerobot_0_3_3 " ;; lerobot) IMAGE="positro/positronic:${IMAGE_TAG}"; EXTRA="--extra lerobot " ;; # openpi.server imports `openpi_client` at module top → needs --extra openpi openpi) IMAGE="positro/openpi:${IMAGE_TAG}"; EXTRA="--extra openpi " ;; - gr00t) IMAGE="positro/gr00t:${IMAGE_TAG}"; EXTRA="" ;; + gr00t) IMAGE="positro/gr00t:${IMAGE_TAG}"; EXTRA="" + UV_ARGS="$GR00T_UV_ARGS"; HF_FLAGS=("${HF_ENV_FLAGS[@]}") ;; # dreamzero.server imports `huggingface_hub` at module top → needs --extra dreamzero dreamzero) IMAGE="positro/dreamzero:${IMAGE_TAG}"; EXTRA="--extra dreamzero " ;; molmoact2) IMAGE="positro/positronic:${IMAGE_TAG}"; EXTRA="--extra molmoact2 " ;; @@ -90,7 +93,8 @@ case " $* " in *) set -- "$@" "--idle_timeout_min=${NEBIUS_IDLE_TIMEOUT_MIN:-20}" ;; esac -SERVER_ARGS="run --python 3.13 ${EXTRA}python -m positronic.vendors.${VENDOR}.server $*" + +SERVER_ARGS="${UV_ARGS} ${EXTRA}python -m positronic.vendors.${VENDOR}.server $*" echo "Creating $VENDOR endpoint '$NAME'..." nebius ai endpoint create \ @@ -109,6 +113,7 @@ nebius ai endpoint create \ --env HF_HOME=/cache/hf \ --env OPENPI_DATA_HOME=/cache/openpi \ --env-secret "${AUTH_TOKEN_KEY}=${AUTH_TOKEN_SECRET}" \ + ${HF_FLAGS[@]+"${HF_FLAGS[@]}"} \ "${S3_ENV_FLAGS[@]}" # Left on stdout, not discarded: `create` reports `Endpoint ID:` as soon as the resource exists and can # still fail afterwards — a container that will not start does exactly that — so a caller logging this diff --git a/workflows/nebius/train.sh b/workflows/nebius/train.sh index ff8fc52e6..2d502f044 100644 --- a/workflows/nebius/train.sh +++ b/workflows/nebius/train.sh @@ -47,11 +47,14 @@ fi VENDOR="$1" shift +UV_ARGS="run --python 3.13" +HF_FLAGS=() case "$VENDOR" in lerobot_0_3_3) IMAGE="positro/positronic:${IMAGE_TAG}"; EXTRA="--extra lerobot_0_3_3 " ;; lerobot) IMAGE="positro/positronic:${IMAGE_TAG}"; EXTRA="--extra lerobot " ;; openpi) IMAGE="positro/openpi:${IMAGE_TAG}"; EXTRA="" ;; - gr00t) IMAGE="positro/gr00t:${IMAGE_TAG}"; EXTRA="" ;; + gr00t) IMAGE="positro/gr00t:${IMAGE_TAG}"; EXTRA="" + UV_ARGS="$GR00T_UV_ARGS"; HF_FLAGS=("${HF_ENV_FLAGS[@]}") ;; dreamzero) IMAGE="positro/dreamzero:${IMAGE_TAG}"; EXTRA="" ;; *) echo "Unknown vendor: '$VENDOR'. Supported: lerobot_0_3_3 | lerobot | openpi | gr00t | dreamzero" >&2 @@ -92,7 +95,8 @@ if [ -n "$INPUT_BUCKET" ]; then fi JOB_NAME="${VENDOR//_/-}-train-$(date +%Y%m%d-%H%M%S)" -TRAIN_ARGS="run --python 3.13 ${EXTRA}python -m positronic.vendors.${VENDOR}.train ${NEW_ARGS[*]}" + +TRAIN_ARGS="${UV_ARGS} ${EXTRA}python -m positronic.vendors.${VENDOR}.train ${NEW_ARGS[*]}" WANDB_FLAGS=() if [ -n "$WANDB_SECRET" ]; then @@ -117,4 +121,5 @@ nebius ai job create \ --env HF_HOME=/cache/hf \ --env OPENPI_DATA_HOME=/cache/openpi \ "${S3_ENV_FLAGS[@]}" \ + ${HF_FLAGS[@]+"${HF_FLAGS[@]}"} \ ${WANDB_FLAGS[@]+"${WANDB_FLAGS[@]}"} From deb2ee15492bd57da9a5bf4b823e82f8bfc769f7 Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Thu, 10 Sep 2026 15:08:14 +0300 Subject: [PATCH 05/14] Pin GR00T to the rebased N1.7 fork commits --- docker/Makefile | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docker/Makefile b/docker/Makefile index 82a27c84d..bd13fc486 100644 --- a/docker/Makefile +++ b/docker/Makefile @@ -43,7 +43,7 @@ TAG_DREAMZERO_BASE_LATEST := $(IMAGE_NAME_DREAMZERO_BASE):latest # OpenPI base image (pulled from registry) TAG_OPENPI_BASE := $(IMAGE_NAME_OPENPI_BASE):latest -GROOT_REF := 506ccc87314fb69ad7ac8d8960b002ce3a378395 +GROOT_REF := 6db6ee8e5702ca2c1c5e1972ff9844dd77d14db2 ifndef GROOT_BASE_IMAGE GROOT_BASE_IMAGE := positro/gr00t-base:$(GROOT_REF) build-groot: build-groot-base From d7c07fc5c4d1b4a428b4325e0c43ff344e353af1 Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Thu, 10 Sep 2026 15:17:28 +0300 Subject: [PATCH 06/14] Update GR00T pin after upstream merge --- docker/Makefile | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docker/Makefile b/docker/Makefile index bd13fc486..d2f59c4d9 100644 --- a/docker/Makefile +++ b/docker/Makefile @@ -43,7 +43,7 @@ TAG_DREAMZERO_BASE_LATEST := $(IMAGE_NAME_DREAMZERO_BASE):latest # OpenPI base image (pulled from registry) TAG_OPENPI_BASE := $(IMAGE_NAME_OPENPI_BASE):latest -GROOT_REF := 6db6ee8e5702ca2c1c5e1972ff9844dd77d14db2 +GROOT_REF := b61c478840e5f79f80ba88ad08f58a09c9d66697 ifndef GROOT_BASE_IMAGE GROOT_BASE_IMAGE := positro/gr00t-base:$(GROOT_REF) build-groot: build-groot-base From f80e384c1e153bdb4e68c1a8c1ff8fee4f5012eb Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Thu, 10 Sep 2026 16:23:18 +0300 Subject: [PATCH 07/14] Validate GR00T camera contracts before serving --- docker/Makefile | 2 +- positronic/policy/codec.py | 5 +++ positronic/policy/observation.py | 2 +- positronic/vendors/dreamzero/codecs.py | 8 ++-- positronic/vendors/gr00t/__init__.py | 13 ++++++ positronic/vendors/gr00t/codecs.py | 44 +++++++++++-------- positronic/vendors/gr00t/server.py | 43 ++++++++++-------- .../vendors/gr00t/tests/test_observation.py | 31 +++++++++++++ positronic/vendors/gr00t/tests/test_server.py | 40 ++++++++++++++++- positronic/vendors/gr00t/train.py | 3 +- .../vendors/lerobot_0_3_3/to_lerobot.py | 11 ++--- positronic/vendors/openpi/codecs.py | 4 +- 12 files changed, 153 insertions(+), 53 deletions(-) diff --git a/docker/Makefile b/docker/Makefile index d2f59c4d9..3cc4689d2 100644 --- a/docker/Makefile +++ b/docker/Makefile @@ -43,7 +43,7 @@ TAG_DREAMZERO_BASE_LATEST := $(IMAGE_NAME_DREAMZERO_BASE):latest # OpenPI base image (pulled from registry) TAG_OPENPI_BASE := $(IMAGE_NAME_OPENPI_BASE):latest -GROOT_REF := b61c478840e5f79f80ba88ad08f58a09c9d66697 +GROOT_REF := 9d17a03fc4d46cc738fd65ca6705eacf3f3c3d6f ifndef GROOT_BASE_IMAGE GROOT_BASE_IMAGE := positro/gr00t-base:$(GROOT_REF) build-groot: build-groot-base diff --git a/positronic/policy/codec.py b/positronic/policy/codec.py index 78c0f6bb2..16eb938fe 100644 --- a/positronic/policy/codec.py +++ b/positronic/policy/codec.py @@ -12,6 +12,7 @@ import collections.abc as cabc from dataclasses import replace from functools import partial +from pathlib import Path from typing import Any, final, overload import numpy as np @@ -29,6 +30,8 @@ from positronic.utils import merge_dicts _QUAT = geom.Rotation.Representation.QUAT +GR00T_MODALITY_PATH = Path('meta/modality.json') +GR00T_MODALITY = 'gr00t_modality' def lerobot_state(dim: int, names: list[str] | None = None) -> dict[str, Any]: @@ -63,6 +66,8 @@ class Codec(Layer): size for all images) or a dict mapping raw input keys to ``(width, height)`` tuples. """ + IMAGE_SIZES = 'image_sizes' + def encode(self, data: dict) -> dict: return {} diff --git a/positronic/policy/observation.py b/positronic/policy/observation.py index ff6aa9b12..192a6561b 100644 --- a/positronic/policy/observation.py +++ b/positronic/policy/observation.py @@ -90,7 +90,7 @@ def encode(self, inputs: dict[str, Any]) -> dict[str, Any]: def meta(self): sizes = {input_key: (w, h) for _out, (input_key, (w, h)) in self._image_configs.items()} unique = set(sizes.values()) - return {'image_sizes': unique.pop() if len(unique) == 1 else sizes} + return {self.IMAGE_SIZES: unique.pop() if len(unique) == 1 else sizes} @property def training_encoder(self): diff --git a/positronic/vendors/dreamzero/codecs.py b/positronic/vendors/dreamzero/codecs.py index 80719cff2..82584efc2 100644 --- a/positronic/vendors/dreamzero/codecs.py +++ b/positronic/vendors/dreamzero/codecs.py @@ -16,7 +16,7 @@ from positronic.drivers.roboarm import command from positronic.drivers.roboarm.ik import DLSIKSolver, DLSIKSolverWithLimits, LMIKSolver from positronic.policy.action import IKJointsAction -from positronic.policy.codec import Codec, lerobot_action, lerobot_image, lerobot_state +from positronic.policy.codec import GR00T_MODALITY, Codec, lerobot_action, lerobot_image, lerobot_state from positronic.vendors.dreamzero import roboarena IMAGE_WIDTH = 320 @@ -66,7 +66,7 @@ def __init__( 'video.exterior_image_1_left': lerobot_image(w, h), 'video.exterior_image_2_left': lerobot_image(w, h), }, - 'gr00t_modality': { + GR00T_MODALITY: { 'state': { 'joint_position': {'start': 0, 'end': 7, 'original_key': 'state.joint_position'}, 'gripper_position': {'start': 0, 'end': 1, 'original_key': 'state.gripper_position'}, @@ -123,7 +123,7 @@ def encode(self, inputs: dict[str, Any]) -> dict[str, Any]: @property def meta(self): - return {'image_sizes': self._image_size} + return {self.IMAGE_SIZES: self._image_size} @property def training_encoder(self): @@ -151,7 +151,7 @@ def __init__(self, tgt_joints_key: str, tgt_grip_key: str, num_joints: int = 7): 'action.joint_position': lerobot_state(num_joints), 'action.gripper_position': lerobot_state(1), }, - 'gr00t_modality': { + GR00T_MODALITY: { 'action': { 'joint_position': {'start': 0, 'end': num_joints, 'original_key': 'action.joint_position'}, 'gripper_position': {'start': 0, 'end': 1, 'original_key': 'action.gripper_position'}, diff --git a/positronic/vendors/gr00t/__init__.py b/positronic/vendors/gr00t/__init__.py index 43e8977d4..4ff7273ce 100644 --- a/positronic/vendors/gr00t/__init__.py +++ b/positronic/vendors/gr00t/__init__.py @@ -1,6 +1,19 @@ """GR00T DROID model and wire vocabulary.""" VIDEO = 'video' +ACTION = 'action' +ANNOTATION = 'annotation' +ORIGINAL_KEY = 'original_key' +TASK_INDEX = 'task_index' +ENDPOINT = 'endpoint' +DATA = 'data' +ERROR = 'error' +OBSERVATION = 'observation' +OPTIONS = 'options' +MODALITY_CONFIG = '__ModalityConfig__' +AS_JSON = 'as_json' +DELTA_INDICES = 'delta_indices' +MODALITY_KEYS = 'modality_keys' STATE = 'state' LANGUAGE = 'language' WRIST_IMAGE = 'wrist_image_left' diff --git a/positronic/vendors/gr00t/codecs.py b/positronic/vendors/gr00t/codecs.py index 8ab076c3e..7ef3d17ae 100644 --- a/positronic/vendors/gr00t/codecs.py +++ b/positronic/vendors/gr00t/codecs.py @@ -15,6 +15,7 @@ from positronic.dataset.transforms.episode import Derive from positronic.drivers.roboarm import command, models from positronic.policy.codec import ( + GR00T_MODALITY, ActionHorizon, ActionTimestamp, BinarizeGripInference, @@ -79,51 +80,58 @@ def _derive_pose(self, episode: Episode): def _derive_grip(episode: Episode): return tf.Elementwise(episode[keys.GRIP], lambda values: np.asarray(values, dtype=np.float32).reshape(-1, 1)) - def _derive_actions(self, episode: Episode): - return tf.concat(self._derive_pose(episode), self._derive_grip(episode), episode[keys.JOINTS], dtype=np.float32) - @staticmethod def _derive_image(source: str, episode: Episode): return image.resize_with_pad(*gr00t.IMAGE_SIZE, signal=episode[source]) @property def training_encoder(self): - state_meta = {name: {'start': 0, 'end': dim, 'original_key': name} for name, dim in gr00t.STATE_DIMS.items()} + state_encoders = { + gr00t.EE_POSE: self._derive_pose, + gr00t.GRIP: self._derive_grip, + gr00t.JOINT_POSITION: lambda episode: tf.Elementwise( + episode[keys.JOINTS], partial(np.asarray, dtype=np.float32) + ), + } + state_meta = { + name: {'start': 0, 'end': gr00t.STATE_DIMS[name], gr00t.ORIGINAL_KEY: name} for name in state_encoders + } action_meta = {} start = 0 - for name, dim in gr00t.STATE_DIMS.items(): + for name in state_encoders: + dim = gr00t.STATE_DIMS[name] action_meta[name] = {'start': start, 'end': start + dim} start += dim meta = { - 'gr00t_modality': { + GR00T_MODALITY: { gr00t.STATE: state_meta, - 'action': action_meta, - gr00t.VIDEO: {name: {'original_key': name} for name in self.image_mappings}, - 'annotation': {gr00t.TASK.removeprefix('annotation.'): {'original_key': 'task_index'}}, + gr00t.ACTION: action_meta, + gr00t.VIDEO: {name: {gr00t.ORIGINAL_KEY: name} for name in self.image_mappings}, + gr00t.ANNOTATION: { + gr00t.TASK.removeprefix(gr00t.ANNOTATION + '.'): {gr00t.ORIGINAL_KEY: gr00t.TASK_INDEX} + }, }, 'lerobot_features': { - **{name: lerobot_state(dim) for name, dim in gr00t.STATE_DIMS.items()}, + **{name: lerobot_state(gr00t.STATE_DIMS[name]) for name in state_encoders}, **{name: lerobot_image(*gr00t.IMAGE_SIZE) for name in self.image_mappings}, - 'action': lerobot_action(start), + gr00t.ACTION: lerobot_action(start), }, } return Derive( meta=meta, **{ - gr00t.EE_POSE: self._derive_pose, - gr00t.GRIP: self._derive_grip, - gr00t.JOINT_POSITION: lambda episode: tf.Elementwise( - episode[keys.JOINTS], partial(np.asarray, dtype=np.float32) - ), + **state_encoders, 'task': itemgetter(keys.TASK), - 'action': self._derive_actions, + gr00t.ACTION: lambda episode: tf.concat( + *(derive(episode) for derive in state_encoders.values()), dtype=np.float32 + ), **{name: partial(self._derive_image, source) for name, source in self.image_mappings.items()}, }, ) @property def meta(self): - return {'image_sizes': dict.fromkeys(self.image_mappings.values(), gr00t.IMAGE_SIZE)} + return {self.IMAGE_SIZES: dict.fromkeys(self.image_mappings.values(), gr00t.IMAGE_SIZE)} @cfn.config( diff --git a/positronic/vendors/gr00t/server.py b/positronic/vendors/gr00t/server.py index c0bdf2028..0989a4aff 100644 --- a/positronic/vendors/gr00t/server.py +++ b/positronic/vendors/gr00t/server.py @@ -18,6 +18,7 @@ from positronic.offboard.server_utils import run_with_progress, wait_for_subprocess_ready, warmup from positronic.policy import Policy, Session from positronic.policy import keys as policy_keys +from positronic.policy.codec import GR00T_MODALITY, Codec, RestrictImageSize from positronic.policy.layers import ChunkedSchedule, StopOnFault from positronic.policy.spec import ModelSource, remote from positronic.utils.checkpoints import list_checkpoints @@ -43,8 +44,8 @@ def decode_custom_classes(obj): if isinstance(obj, dict): if obj.get(b'nd', obj.get('nd')) and obj.get(b'kind', obj.get('kind')) in (b'O', 'O'): raise ValueError('Object arrays are not supported by the GR00T wire protocol') - if obj.get('__ModalityConfig__'): - return obj['as_json'] + if obj.get(gr00t.MODALITY_CONFIG): + return obj[gr00t.AS_JSON] return mnp.decode(obj) @staticmethod @@ -79,9 +80,9 @@ def ping(self) -> bool: return False def call_endpoint(self, endpoint: str, data: dict | None = None) -> Any: - request: dict = {'endpoint': endpoint} + request: dict = {gr00t.ENDPOINT: endpoint} if data is not None: - request['data'] = data + request[gr00t.DATA] = data try: self.socket.send(MsgSerializer.to_bytes(request)) @@ -95,16 +96,16 @@ def call_endpoint(self, endpoint: str, data: dict | None = None) -> Any: raise RuntimeError('Server error. Make sure the correct policy server is running.') response = MsgSerializer.from_bytes(message) - if isinstance(response, dict) and 'error' in response: - raise RuntimeError(f'Server error: {response["error"]}') + if isinstance(response, dict) and gr00t.ERROR in response: + raise RuntimeError(f'Server error: {response[gr00t.ERROR]}') return response def get_action(self, observation: dict[str, Any]) -> tuple[dict, dict]: - response = self.call_endpoint('get_action', {'observation': observation, 'options': None}) + response = self.call_endpoint('get_action', {gr00t.OBSERVATION: observation, gr00t.OPTIONS: None}) return tuple(response) def reset(self) -> dict[str, Any]: - return self.call_endpoint('reset', {'options': None}) + return self.call_endpoint('reset', {gr00t.OPTIONS: None}) def close(self): self.socket.close(linger=0) @@ -216,12 +217,14 @@ class Gr00tSource(ModelSource): def __init__( self, + video_keys: tuple[str, ...], checkpoints_dir: str = 'hf://' + gr00t.BASE_MODEL, checkpoint: str | None = None, groot_venv_path: str = gr00t.VENV, zmq_port: int = 5555, ready_timeout: float = 600.0, ): + self.video_keys = tuple(video_keys) self.checkpoints_dir = checkpoints_dir.rstrip('/') if self._is_hub_model and checkpoint is not None: raise ValueError('checkpoint step selection applies only to fine-tuned checkpoint directories') @@ -271,23 +274,24 @@ def resolve(self, model_id: str | None) -> str: return self._step_id(self._raw_ids()[-1]) return self._step_id(self._raw_for(model_id)) - @staticmethod - def _warm_observation(modalities: dict) -> dict[str, Any]: - """Build a valid current-frame DROID observation from the loaded checkpoint's modalities.""" + def _warm_observation(self, modalities: dict) -> dict[str, Any]: + """Validate codec cameras and build a current-frame observation for the checkpoint.""" for name in (gr00t.VIDEO, gr00t.STATE): - if modalities[name]['delta_indices'] != [0]: + if modalities[name][gr00t.DELTA_INDICES] != [0]: raise ValueError(f'DROID adapter requires current-frame {name}, got {modalities[name]}') + expected = set(modalities[gr00t.VIDEO][gr00t.MODALITY_KEYS]) + if set(self.video_keys) != expected: + raise ValueError( + f'Checkpoint video keys {sorted(expected)} do not match codec keys {sorted(self.video_keys)}' + ) width, height = gr00t.IMAGE_SIZE state = { name: np.zeros((1, 1, gr00t.STATE_DIMS[name]), dtype=np.float32) - for name in modalities[gr00t.STATE]['modality_keys'] + for name in modalities[gr00t.STATE][gr00t.MODALITY_KEYS] } state[gr00t.EE_POSE][..., 3:] = [1, 0, 0, 0, 1, 0] return { - gr00t.VIDEO: { - name: np.zeros((1, 1, height, width, 3), dtype=np.uint8) - for name in modalities[gr00t.VIDEO]['modality_keys'] - }, + gr00t.VIDEO: {name: np.zeros((1, 1, height, width, 3), dtype=np.uint8) for name in self.video_keys}, gr00t.STATE: state, gr00t.LANGUAGE: {gr00t.TASK: [['pick up the object']]}, } @@ -332,9 +336,10 @@ def meta(self, model_id: str) -> dict[str, Any]: @cfn.config(codec=codecs.droid, source=gr00t_source) -def pipeline(codec, source): +def pipeline(codec: Codec, source: cfn.Config): """Schedule DROID joint commands while the server codec performs checkpoint-specific conversion.""" - return StopOnFault() | ChunkedSchedule() | remote | codec | source + model_source = source(video_keys=tuple(codec.training_encoder.meta[GR00T_MODALITY][gr00t.VIDEO])) + return StopOnFault() | ChunkedSchedule() | RestrictImageSize(*gr00t.IMAGE_SIZE) | remote | codec | model_source droid = pipeline diff --git a/positronic/vendors/gr00t/tests/test_observation.py b/positronic/vendors/gr00t/tests/test_observation.py index ef8773b05..2e1cba5a9 100644 --- a/positronic/vendors/gr00t/tests/test_observation.py +++ b/positronic/vendors/gr00t/tests/test_observation.py @@ -10,7 +10,10 @@ from positronic.dataset.episode import EpisodeContainer from positronic.dataset.tests.utils import DummySignal from positronic.drivers.roboarm import models +from positronic.policy.codec import GR00T_MODALITY, Codec, RestrictImageSize +from positronic.policy.spec import split from positronic.vendors import gr00t +from positronic.vendors.gr00t import server from positronic.vendors.gr00t.codecs import droid, droid_three_cameras @@ -57,6 +60,34 @@ def test_three_camera_configuration_uses_a_distinct_second_external_image(observ droid_three_cameras().encode(observation) +def test_action_metadata_matches_values_when_state_dimensions_are_reordered(monkeypatch, observation): + monkeypatch.setattr(gr00t, 'STATE_DIMS', dict(reversed(list(gr00t.STATE_DIMS.items())))) + codec = droid() + episode = EpisodeContainer({ + name: value if name == keys.TASK else DummySignal([0, 1], [value, value]) for name, value in observation.items() + }) + encoder = codec.training_encoder + encoded = encoder(episode) + action = encoded[gr00t.ACTION][0][0] + for name, bounds in encoder.meta[GR00T_MODALITY][gr00t.ACTION].items(): + np.testing.assert_allclose(action[bounds['start'] : bounds['end']], encoded[name][0][0]) + + +@pytest.mark.parametrize('config', [server.droid, server.droid_three_cameras]) +def test_images_are_bounded_before_remote_without_changing_model_pixels(config, observation): + pipeline = config() + local, _, codec = split(pipeline) + resize = next(layer for layer in local._layers() if isinstance(layer, RestrictImageSize)) + wire_observation = resize.encode(observation) + for source in codec.meta[Codec.IMAGE_SIZES]: + assert wire_observation[source].shape[0] <= gr00t.IMAGE_SIZE[1] + assert wire_observation[source].shape[1] <= gr00t.IMAGE_SIZE[0] + direct = codec.encode(observation) + remote_encoded = codec.encode(wire_observation) + for name in direct[gr00t.VIDEO]: + np.testing.assert_array_equal(remote_encoded[gr00t.VIDEO][name], direct[gr00t.VIDEO][name]) + + def test_droid_frame_and_pixels_match_upstream_robot_client(observation): reference = os.environ.get('GR00T_REFERENCE_ROOT') if reference is None: diff --git a/positronic/vendors/gr00t/tests/test_server.py b/positronic/vendors/gr00t/tests/test_server.py index 72ccff287..3d073abdd 100644 --- a/positronic/vendors/gr00t/tests/test_server.py +++ b/positronic/vendors/gr00t/tests/test_server.py @@ -6,13 +6,14 @@ import pytest import zmq +from positronic.offboard.server import PolicyServer from positronic.vendors import gr00t from positronic.vendors.gr00t import server as gr00t_server def _source(monkeypatch, checkpoints: list[str]) -> gr00t_server.Gr00tSource: monkeypatch.setattr(gr00t_server, 'list_checkpoints', lambda _dir, prefix='': checkpoints) - return gr00t_server.Gr00tSource('s3://bucket/exp') + return gr00t_server.droid.override_data(**{'source.checkpoints_dir': 's3://bucket/exp'})().source def test_zero_padded_checkpoints_are_served_under_the_id_they_advertise(monkeypatch): @@ -49,11 +50,46 @@ def test_serializer_rejects_pickle_bearing_arrays(): def test_published_checkpoint_is_served_without_a_local_checkpoint_scan(monkeypatch): - source = gr00t_server.Gr00tSource() + source = gr00t_server.droid().source assert source.get_models() == [gr00t.BASE_MODEL] assert source.resolve(None) == gr00t.BASE_MODEL +@pytest.mark.parametrize('config', [gr00t_server.droid, gr00t_server.droid_three_cameras]) +@pytest.mark.parametrize('checkpoint_cameras', [2, 3]) +def test_camera_mismatch_stops_the_backend_before_warmup(monkeypatch, config, checkpoint_cameras): + source = config().source + cameras = [gr00t.EXTERIOR_IMAGE, gr00t.WRIST_IMAGE] + if checkpoint_cameras == 3: + cameras.append(gr00t.EXTERIOR_IMAGE_2) + backend = Mock() + backend.client.call_endpoint.return_value = { + gr00t.VIDEO: {gr00t.DELTA_INDICES: [0], gr00t.MODALITY_KEYS: cameras}, + gr00t.STATE: {gr00t.DELTA_INDICES: [0], gr00t.MODALITY_KEYS: list(gr00t.STATE_DIMS)}, + } + monkeypatch.setattr(gr00t_server, 'Gr00tSubprocess', Mock(return_value=backend)) + warmup = Mock() + monkeypatch.setattr(gr00t_server, 'warmup', warmup) + if len(source.video_keys) != checkpoint_cameras: + with pytest.raises(ValueError, match='Checkpoint video keys'): + source.load(gr00t.BASE_MODEL) + warmup.assert_not_called() + backend.stop.assert_called_once() + else: + policy = source.load(gr00t.BASE_MODEL) + try: + assert set(warmup.call_args.args[1][gr00t.VIDEO]) == set(cameras) + backend.stop.assert_not_called() + finally: + policy.close() + + +def test_session_timing_overrides_preserve_source_equality(): + server = PolicyServer(gr00t_server.droid) + variant = server._session_pipeline({'codec.fps': 10.0}) + assert variant.source == gr00t_server.droid().source + + @pytest.mark.parametrize('failure', [zmq.Again(), zmq.ZMQError(zmq.EFSM)]) def test_client_can_ping_after_a_transport_failure(monkeypatch, failure): failed = Mock() diff --git a/positronic/vendors/gr00t/train.py b/positronic/vendors/gr00t/train.py index f1129a91c..270845672 100644 --- a/positronic/vendors/gr00t/train.py +++ b/positronic/vendors/gr00t/train.py @@ -7,6 +7,7 @@ import pos3 from positronic import utils +from positronic.policy.codec import GR00T_MODALITY_PATH from positronic.vendors import gr00t @@ -41,7 +42,7 @@ def main( python_bin = str(Path(groot_venv_path).expanduser() / 'bin' / 'python') dataset_local_path = pos3.download(input_path) - with (Path(dataset_local_path) / 'meta' / 'modality.json').open() as f: + with (Path(dataset_local_path) / GR00T_MODALITY_PATH).open() as f: video_keys = list(json.load(f)[gr00t.VIDEO]) output_path = output_path.rstrip('/') output_dir = pos3.sync(output_path + '/' + exp_name, delete_remote=not resume) diff --git a/positronic/vendors/lerobot_0_3_3/to_lerobot.py b/positronic/vendors/lerobot_0_3_3/to_lerobot.py index 3d978e607..928e9e242 100644 --- a/positronic/vendors/lerobot_0_3_3/to_lerobot.py +++ b/positronic/vendors/lerobot_0_3_3/to_lerobot.py @@ -31,6 +31,7 @@ from positronic import utils from positronic.cfg.ds import apply_codec from positronic.dataset import Dataset +from positronic.policy.codec import GR00T_MODALITY, GR00T_MODALITY_PATH def _raise_fd_limit(min_soft_limit: int = 4096) -> None: @@ -142,10 +143,10 @@ def convert_to_lerobot_dataset( # otherwise the former will complain about the directory not being empty. utils.save_run_metadata(output_dir, patterns=['*.py', '*.toml']) - if 'gr00t_modality' in dataset.meta: - modality = dataset.meta.get('gr00t_modality') + if GR00T_MODALITY in dataset.meta: + modality = dataset.meta.get(GR00T_MODALITY) if modality is not None: - modality_path = output_dir / 'meta' / 'modality.json' + modality_path = output_dir / GR00T_MODALITY_PATH with modality_path.open('w', encoding='utf-8') as f: json.dump(modality, f, indent=2) @@ -164,8 +165,8 @@ def append_data_to_lerobot_dataset(output_dir: str, dataset: Dataset, fps: int | # Save metadata for append operation utils.save_run_metadata(output_dir, patterns=['*.py', '*.toml'], prefix='append_metadata') - lr_modality_path = output_dir / 'meta' / 'modality.json' - ds_modality = dataset.meta.get('gr00t_modality', None) + lr_modality_path = output_dir / GR00T_MODALITY_PATH + ds_modality = dataset.meta.get(GR00T_MODALITY, None) if lr_modality_path.exists(): with lr_modality_path.open(encoding='utf-8') as f: lr_modality = json.load(f) diff --git a/positronic/vendors/openpi/codecs.py b/positronic/vendors/openpi/codecs.py index 5979b4f60..7f6f3f95c 100644 --- a/positronic/vendors/openpi/codecs.py +++ b/positronic/vendors/openpi/codecs.py @@ -103,7 +103,7 @@ def _encode_image(self, input_key: str, inputs: dict[str, Any]) -> np.ndarray: @property def meta(self): - return {'image_sizes': self._image_size} + return {self.IMAGE_SIZES: self._image_size} @property def training_encoder(self): @@ -264,7 +264,7 @@ def _encode_image(self, input_key: str, inputs: dict[str, Any]) -> np.ndarray: @property def meta(self) -> dict[str, Any]: - return {'image_sizes': self._image_size} + return {self.IMAGE_SIZES: self._image_size} libero_obs = cfn.Config(LiberoObservationCodec) From 6302abd7cdf1ea7447bb3d53a7792915edd2900c Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Thu, 10 Sep 2026 16:51:04 +0300 Subject: [PATCH 08/14] Share dataset metadata fields and pin reviewed GR00T --- docker/Makefile | 2 +- positronic/policy/codec.py | 1 + positronic/vendors/dreamzero/codecs.py | 15 +++++++++++---- positronic/vendors/gr00t/README.md | 7 +++++++ positronic/vendors/gr00t/__init__.py | 2 ++ positronic/vendors/gr00t/codecs.py | 10 ++++++---- positronic/vendors/lerobot/to_lerobot.py | 11 ++++++----- positronic/vendors/lerobot/train.py | 7 ++++--- positronic/vendors/lerobot_0_3_3/to_lerobot.py | 12 ++++++------ positronic/vendors/lerobot_0_3_3/train.py | 7 ++++--- positronic/vendors/openpi/codecs.py | 6 +++--- 11 files changed, 51 insertions(+), 29 deletions(-) diff --git a/docker/Makefile b/docker/Makefile index 3cc4689d2..e40e9ba9c 100644 --- a/docker/Makefile +++ b/docker/Makefile @@ -43,7 +43,7 @@ TAG_DREAMZERO_BASE_LATEST := $(IMAGE_NAME_DREAMZERO_BASE):latest # OpenPI base image (pulled from registry) TAG_OPENPI_BASE := $(IMAGE_NAME_OPENPI_BASE):latest -GROOT_REF := 9d17a03fc4d46cc738fd65ca6705eacf3f3c3d6f +GROOT_REF := e63c6f70257574cd7cf41a5abd24d3f0cab1306c ifndef GROOT_BASE_IMAGE GROOT_BASE_IMAGE := positro/gr00t-base:$(GROOT_REF) build-groot: build-groot-base diff --git a/positronic/policy/codec.py b/positronic/policy/codec.py index 16eb938fe..bc4754d0f 100644 --- a/positronic/policy/codec.py +++ b/positronic/policy/codec.py @@ -32,6 +32,7 @@ _QUAT = geom.Rotation.Representation.QUAT GR00T_MODALITY_PATH = Path('meta/modality.json') GR00T_MODALITY = 'gr00t_modality' +LEROBOT_FEATURES = 'lerobot_features' def lerobot_state(dim: int, names: list[str] | None = None) -> dict[str, Any]: diff --git a/positronic/vendors/dreamzero/codecs.py b/positronic/vendors/dreamzero/codecs.py index 82584efc2..2f1a0b1a0 100644 --- a/positronic/vendors/dreamzero/codecs.py +++ b/positronic/vendors/dreamzero/codecs.py @@ -16,7 +16,14 @@ from positronic.drivers.roboarm import command from positronic.drivers.roboarm.ik import DLSIKSolver, DLSIKSolverWithLimits, LMIKSolver from positronic.policy.action import IKJointsAction -from positronic.policy.codec import GR00T_MODALITY, Codec, lerobot_action, lerobot_image, lerobot_state +from positronic.policy.codec import ( + GR00T_MODALITY, + LEROBOT_FEATURES, + Codec, + lerobot_action, + lerobot_image, + lerobot_state, +) from positronic.vendors.dreamzero import roboarena IMAGE_WIDTH = 320 @@ -55,11 +62,11 @@ def __init__( 'video.wrist_image_left': partial(self._derive_image, wrist_camera), 'video.exterior_image_1_left': partial(self._derive_image, exterior_camera_1), 'video.exterior_image_2_left': partial(self._derive_image, self._exterior_camera_2), - 'task': Get(keys.TASK, ''), + keys.TASK: Get(keys.TASK, ''), } self._training_meta = { - 'lerobot_features': { + LEROBOT_FEATURES: { 'state.joint_position': lerobot_state(7), 'state.gripper_position': lerobot_state(1), 'video.wrist_image_left': lerobot_image(w, h), @@ -146,7 +153,7 @@ def __init__(self, tgt_joints_key: str, tgt_grip_key: str, num_joints: int = 7): self._num_joints = num_joints self._training_meta = { - 'lerobot_features': { + LEROBOT_FEATURES: { 'action': lerobot_action(num_joints + 1), 'action.joint_position': lerobot_state(num_joints), 'action.gripper_position': lerobot_state(1), diff --git a/positronic/vendors/gr00t/README.md b/positronic/vendors/gr00t/README.md index 2509cd032..0755f3ab9 100644 --- a/positronic/vendors/gr00t/README.md +++ b/positronic/vendors/gr00t/README.md @@ -25,6 +25,13 @@ The published two-camera checkpoint does not consume an extra exterior view. Sel `droid_three_cameras` for both conversion and serving when fine-tuning with three views. N1.6 checkpoints require an N1.6 image; their custom action schemas are incompatible with this adapter. +The default pose transform matches Franka rigs and the RoboLab adapter. The bundled MuJoCo +Panda reports poses at a different tool frame; using the default transform on its recordings +(including `sim_stack_cubes`) introduces a 45 mm offset. That frame alignment is tracked in +[#550](https://github.com/Positronic-Robotics/positronic/issues/550). Such datasets and their +serving codec require a matching simulator-specific `ee_frame`; the default is not a validated +native-checkpoint configuration for that simulator. + ## Docker Build the pinned fork and the Positronic image: diff --git a/positronic/vendors/gr00t/__init__.py b/positronic/vendors/gr00t/__init__.py index 4ff7273ce..23498f94d 100644 --- a/positronic/vendors/gr00t/__init__.py +++ b/positronic/vendors/gr00t/__init__.py @@ -4,6 +4,8 @@ ACTION = 'action' ANNOTATION = 'annotation' ORIGINAL_KEY = 'original_key' +START = 'start' +END = 'end' TASK_INDEX = 'task_index' ENDPOINT = 'endpoint' DATA = 'data' diff --git a/positronic/vendors/gr00t/codecs.py b/positronic/vendors/gr00t/codecs.py index 7ef3d17ae..7fafdf48e 100644 --- a/positronic/vendors/gr00t/codecs.py +++ b/positronic/vendors/gr00t/codecs.py @@ -16,6 +16,7 @@ from positronic.drivers.roboarm import command, models from positronic.policy.codec import ( GR00T_MODALITY, + LEROBOT_FEATURES, ActionHorizon, ActionTimestamp, BinarizeGripInference, @@ -94,13 +95,14 @@ def training_encoder(self): ), } state_meta = { - name: {'start': 0, 'end': gr00t.STATE_DIMS[name], gr00t.ORIGINAL_KEY: name} for name in state_encoders + name: {gr00t.START: 0, gr00t.END: gr00t.STATE_DIMS[name], gr00t.ORIGINAL_KEY: name} + for name in state_encoders } action_meta = {} start = 0 for name in state_encoders: dim = gr00t.STATE_DIMS[name] - action_meta[name] = {'start': start, 'end': start + dim} + action_meta[name] = {gr00t.START: start, gr00t.END: start + dim} start += dim meta = { GR00T_MODALITY: { @@ -111,7 +113,7 @@ def training_encoder(self): gr00t.TASK.removeprefix(gr00t.ANNOTATION + '.'): {gr00t.ORIGINAL_KEY: gr00t.TASK_INDEX} }, }, - 'lerobot_features': { + LEROBOT_FEATURES: { **{name: lerobot_state(gr00t.STATE_DIMS[name]) for name in state_encoders}, **{name: lerobot_image(*gr00t.IMAGE_SIZE) for name in self.image_mappings}, gr00t.ACTION: lerobot_action(start), @@ -121,7 +123,7 @@ def training_encoder(self): meta=meta, **{ **state_encoders, - 'task': itemgetter(keys.TASK), + keys.TASK: itemgetter(keys.TASK), gr00t.ACTION: lambda episode: tf.concat( *(derive(episode) for derive in state_encoders.values()), dtype=np.float32 ), diff --git a/positronic/vendors/lerobot/to_lerobot.py b/positronic/vendors/lerobot/to_lerobot.py index c046fdb50..f003e6685 100644 --- a/positronic/vendors/lerobot/to_lerobot.py +++ b/positronic/vendors/lerobot/to_lerobot.py @@ -25,9 +25,10 @@ from lerobot.datasets.lerobot_dataset import LeRobotDataset from pimm.logging import init_logging -from positronic import utils +from positronic import keys, utils from positronic.cfg.ds import apply_codec from positronic.dataset import Dataset +from positronic.policy.codec import LEROBOT_FEATURES def _raise_fd_limit(min_soft_limit: int = 4096) -> None: @@ -98,8 +99,8 @@ def append_data_to_dataset( value = value[i] frame[key] = value - ep_task = task if task is not None else frame.get('task', '') - frame['task'] = ep_task or '' + ep_task = task if task is not None else frame.get(keys.TASK, '') + frame[keys.TASK] = ep_task or '' lr_dataset.add_frame(frame) lr_dataset.save_episode(parallel_encoding=False) @@ -115,14 +116,14 @@ def convert_to_lerobot_dataset( assert 'action_fps' in dataset.meta, "--fps not provided and dataset has no 'action_fps' metadata" fps = int(dataset.meta['action_fps']) output_dir = pos3.sync(output_dir, interval=None, sync_on_error=False) - assert dataset.meta['lerobot_features'] is not None, "dataset.meta['lerobot_features'] is required" + assert dataset.meta[LEROBOT_FEATURES] is not None, f'dataset.meta[{LEROBOT_FEATURES!r}] is required' lr_dataset = LeRobotDataset.create( repo_id='local', fps=fps, root=output_dir, use_videos=video, - features=dataset.meta['lerobot_features'], + features=dataset.meta[LEROBOT_FEATURES], image_writer_threads=32, ) utils.save_run_metadata(output_dir, patterns=['*.py', '*.toml']) diff --git a/positronic/vendors/lerobot/train.py b/positronic/vendors/lerobot/train.py index 3a0b8aa82..e7c9d7665 100644 --- a/positronic/vendors/lerobot/train.py +++ b/positronic/vendors/lerobot/train.py @@ -37,6 +37,7 @@ from pimm.logging import init_logging from positronic import utils from positronic.policy import Codec +from positronic.policy.codec import LEROBOT_FEATURES from positronic.vendors.lerobot import codecs as lerobot_codecs # Workaround: HubMixin.save_pretrained unconditionally deletes config.json before @@ -77,10 +78,10 @@ def build_env_config_from_codec(codec: Codec) -> PositronicEnvConfig: fps = int(inference_meta.get('action_fps', 15)) - assert 'lerobot_features' in training_meta, ( - f"Codec training_encoder missing 'lerobot_features'. Keys: {list(training_meta.keys())}" + assert LEROBOT_FEATURES in training_meta, ( + f'Codec training_encoder missing {LEROBOT_FEATURES!r}. Keys: {list(training_meta.keys())}' ) - lerobot_features = training_meta['lerobot_features'] + lerobot_features = training_meta[LEROBOT_FEATURES] features = {} features_map = {} diff --git a/positronic/vendors/lerobot_0_3_3/to_lerobot.py b/positronic/vendors/lerobot_0_3_3/to_lerobot.py index 928e9e242..ceb0b113c 100644 --- a/positronic/vendors/lerobot_0_3_3/to_lerobot.py +++ b/positronic/vendors/lerobot_0_3_3/to_lerobot.py @@ -28,10 +28,10 @@ from lerobot.datasets.lerobot_dataset import LeRobotDataset from pimm.logging import init_logging -from positronic import utils +from positronic import keys, utils from positronic.cfg.ds import apply_codec from positronic.dataset import Dataset -from positronic.policy.codec import GR00T_MODALITY, GR00T_MODALITY_PATH +from positronic.policy.codec import GR00T_MODALITY, GR00T_MODALITY_PATH, LEROBOT_FEATURES def _raise_fd_limit(min_soft_limit: int = 4096) -> None: @@ -111,9 +111,9 @@ def append_data_to_dataset( ep_task = task if task is None: - ep_task = frame.get('task', '') + ep_task = frame.get(keys.TASK, '') - frame.pop('task', None) + frame.pop(keys.TASK, None) lr_dataset.add_frame(frame, task=ep_task or '') lr_dataset.save_episode() @@ -129,14 +129,14 @@ def convert_to_lerobot_dataset( assert 'action_fps' in dataset.meta, "--fps not provided and dataset has no 'action_fps' metadata" fps = int(dataset.meta['action_fps']) output_dir = pos3.sync(output_dir, interval=None, sync_on_error=False) - assert dataset.meta['lerobot_features'] is not None, "dataset.meta['lerobot_features'] is required" + assert dataset.meta[LEROBOT_FEATURES] is not None, f'dataset.meta[{LEROBOT_FEATURES!r}] is required' lr_dataset = LeRobotDataset.create( repo_id='local', fps=fps, root=output_dir, use_videos=video, - features=dataset.meta['lerobot_features'], + features=dataset.meta[LEROBOT_FEATURES], image_writer_threads=32, ) # Adding this file after the LR dataset is created, diff --git a/positronic/vendors/lerobot_0_3_3/train.py b/positronic/vendors/lerobot_0_3_3/train.py index 6f0f2e0de..ae42433cf 100644 --- a/positronic/vendors/lerobot_0_3_3/train.py +++ b/positronic/vendors/lerobot_0_3_3/train.py @@ -32,6 +32,7 @@ from pimm.logging import init_logging from positronic import utils from positronic.policy import Codec +from positronic.policy.codec import LEROBOT_FEATURES from positronic.vendors.lerobot_0_3_3 import codecs as lerobot_codecs from positronic.vendors.lerobot_0_3_3.backbone import BACKBONES @@ -57,10 +58,10 @@ def build_env_config_from_codec(codec: Codec) -> PositronicEnvConfig: fps = int(inference_meta.get('action_fps', 15)) - assert 'lerobot_features' in training_meta, ( - f"Codec training_encoder missing 'lerobot_features'. Keys: {list(training_meta.keys())}" + assert LEROBOT_FEATURES in training_meta, ( + f'Codec training_encoder missing {LEROBOT_FEATURES!r}. Keys: {list(training_meta.keys())}' ) - lerobot_features = training_meta['lerobot_features'] + lerobot_features = training_meta[LEROBOT_FEATURES] features = {} features_map = {} diff --git a/positronic/vendors/openpi/codecs.py b/positronic/vendors/openpi/codecs.py index 7f6f3f95c..c714469ee 100644 --- a/positronic/vendors/openpi/codecs.py +++ b/positronic/vendors/openpi/codecs.py @@ -29,7 +29,7 @@ from positronic.dataset.transforms import image from positronic.dataset.transforms.episode import Derive, Get from positronic.drivers.roboarm import command -from positronic.policy.codec import Codec, lerobot_image, lerobot_state +from positronic.policy.codec import LEROBOT_FEATURES, Codec, lerobot_image, lerobot_state from positronic.policy.observation import ObservationCodec as GenericObservationCodec from positronic.vendors import openpi @@ -53,13 +53,13 @@ def __init__( 'observation.state': self._derive_state, 'observation.images.left': partial(self._derive_image, wrist_camera), 'observation.images.side': partial(self._derive_image, exterior_camera), - 'task': Get(keys.TASK, ''), + keys.TASK: Get(keys.TASK, ''), } state_dim = sum(state_features.values()) w, h = image_size self._training_meta: dict[str, Any] = { - 'lerobot_features': { + LEROBOT_FEATURES: { 'observation.state': lerobot_state(state_dim, list(state_features.keys())), 'observation.images.left': lerobot_image(w, h), 'observation.images.side': lerobot_image(w, h), From f8f82e9c5f7344b7339ca1ce07c695ce2d396434 Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Thu, 10 Sep 2026 17:08:52 +0300 Subject: [PATCH 09/14] Name GR00T client endpoints in the wire vocabulary --- positronic/vendors/gr00t/__init__.py | 4 ++++ positronic/vendors/gr00t/server.py | 8 ++++---- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/positronic/vendors/gr00t/__init__.py b/positronic/vendors/gr00t/__init__.py index 23498f94d..960f1c10a 100644 --- a/positronic/vendors/gr00t/__init__.py +++ b/positronic/vendors/gr00t/__init__.py @@ -8,6 +8,10 @@ END = 'end' TASK_INDEX = 'task_index' ENDPOINT = 'endpoint' +PING = 'ping' +GET_ACTION = 'get_action' +RESET = 'reset' +GET_MODALITY_CONFIG = 'get_modality_config' DATA = 'data' ERROR = 'error' OBSERVATION = 'observation' diff --git a/positronic/vendors/gr00t/server.py b/positronic/vendors/gr00t/server.py index 0989a4aff..b6c579a95 100644 --- a/positronic/vendors/gr00t/server.py +++ b/positronic/vendors/gr00t/server.py @@ -74,7 +74,7 @@ def _make_socket(self): def ping(self) -> bool: try: - self.call_endpoint('ping') + self.call_endpoint(gr00t.PING) return True except (zmq.error.ZMQError, RuntimeError): return False @@ -101,11 +101,11 @@ def call_endpoint(self, endpoint: str, data: dict | None = None) -> Any: return response def get_action(self, observation: dict[str, Any]) -> tuple[dict, dict]: - response = self.call_endpoint('get_action', {gr00t.OBSERVATION: observation, gr00t.OPTIONS: None}) + response = self.call_endpoint(gr00t.GET_ACTION, {gr00t.OBSERVATION: observation, gr00t.OPTIONS: None}) return tuple(response) def reset(self) -> dict[str, Any]: - return self.call_endpoint('reset', {gr00t.OPTIONS: None}) + return self.call_endpoint(gr00t.RESET, {gr00t.OPTIONS: None}) def close(self): self.socket.close(linger=0) @@ -317,7 +317,7 @@ def load(self, model_id: str, on_progress: Callable[[str], None] | None = None) groot.start(on_progress) policy = Gr00tPolicy(groot, str(checkpoint_dir)) # The subprocess initializes CUDA on its first forward, which outlasts a rig's inference timeout. - modalities = groot.client.call_endpoint('get_modality_config') + modalities = groot.client.call_endpoint(gr00t.GET_MODALITY_CONFIG) warmup(policy, self._warm_observation(modalities), on_progress) except Exception: groot.stop() From 7e36bbd39f7bfbf5b6cc4a05d0dfd0a1ec06dc1b Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Thu, 10 Sep 2026 17:27:30 +0300 Subject: [PATCH 10/14] Keep desktop PhAIL on the N1.6 server image --- docker/docker-compose.phail.desktop.yml | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/docker/docker-compose.phail.desktop.yml b/docker/docker-compose.phail.desktop.yml index 2a0c53348..e91b7d2da 100644 --- a/docker/docker-compose.phail.desktop.yml +++ b/docker/docker-compose.phail.desktop.yml @@ -1,13 +1,11 @@ -# GR00T on desktop (RTX 3060 12GB — GR00T uses ~4.6GB) +# GR00T N1.6 on desktop (RTX 3060 12GB — GR00T uses ~4.6GB) # +# Set GR00T_N16_IMAGE to a pinned N1.6 image. # docker --context desktop compose -f docker-compose.phail.desktop.yml up services: phail-groot-server: extends: - file: docker-compose.yml - service: groot-server - container_name: phail-groot-server - pull_policy: always - command: ["phail"] + file: docker-compose.phail.yml + service: phail-groot-server ports: !override - "8000:8000" From 5a19de90d68b94c27eec4d6f6cc6bb91fe8eea73 Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Thu, 10 Sep 2026 18:29:44 +0300 Subject: [PATCH 11/14] Restore dependency sync for legacy GR00T images --- docker/docker-compose.phail.yml | 2 +- docker/docker-compose.spoons-ablation.yml | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/docker/docker-compose.phail.yml b/docker/docker-compose.phail.yml index e97fd39ed..0d08edbd0 100644 --- a/docker/docker-compose.phail.yml +++ b/docker/docker-compose.phail.yml @@ -15,7 +15,7 @@ services: service: groot-server # These experiment checkpoints use the N1.6 action schema. image: ${GR00T_N16_IMAGE:?Set GR00T_N16_IMAGE to a pinned GR00T N1.6 image} - entrypoint: ["uv", "run", "--no-sync", "--python", "3.13", "python", "-m", "positronic.vendors.gr00t.server"] + entrypoint: ["uv", "run", "--python", "3.13", "python", "-m", "positronic.vendors.gr00t.server"] container_name: phail-groot-server pull_policy: always command: ["phail"] diff --git a/docker/docker-compose.spoons-ablation.yml b/docker/docker-compose.spoons-ablation.yml index ec0fb2573..a4553858d 100644 --- a/docker/docker-compose.spoons-ablation.yml +++ b/docker/docker-compose.spoons-ablation.yml @@ -11,7 +11,7 @@ services: service: groot-server # These experiment checkpoints use the N1.6 action schema. image: ${GR00T_N16_IMAGE:?Set GR00T_N16_IMAGE to a pinned GR00T N1.6 image} - entrypoint: ["uv", "run", "--no-sync", "--python", "3.13", "python", "-m", "positronic.vendors.gr00t.server"] + entrypoint: ["uv", "run", "--python", "3.13", "python", "-m", "positronic.vendors.gr00t.server"] container_name: spoons-100 pull_policy: always command: @@ -26,7 +26,7 @@ services: service: groot-server # These experiment checkpoints use the N1.6 action schema. image: ${GR00T_N16_IMAGE:?Set GR00T_N16_IMAGE to a pinned GR00T N1.6 image} - entrypoint: ["uv", "run", "--no-sync", "--python", "3.13", "python", "-m", "positronic.vendors.gr00t.server"] + entrypoint: ["uv", "run", "--python", "3.13", "python", "-m", "positronic.vendors.gr00t.server"] container_name: spoons-50 pull_policy: always command: @@ -41,7 +41,7 @@ services: service: groot-server # These experiment checkpoints use the N1.6 action schema. image: ${GR00T_N16_IMAGE:?Set GR00T_N16_IMAGE to a pinned GR00T N1.6 image} - entrypoint: ["uv", "run", "--no-sync", "--python", "3.13", "python", "-m", "positronic.vendors.gr00t.server"] + entrypoint: ["uv", "run", "--python", "3.13", "python", "-m", "positronic.vendors.gr00t.server"] container_name: spoons-25 pull_policy: always command: From b03a2abf56e010078df03c04b71d13ae0dfa0004 Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Thu, 10 Sep 2026 18:44:49 +0300 Subject: [PATCH 12/14] Name GR00T wire descriptors and ping outcomes --- positronic/vendors/gr00t/server.py | 22 ++++++++++++++----- positronic/vendors/gr00t/tests/test_server.py | 4 ++-- 2 files changed, 19 insertions(+), 7 deletions(-) diff --git a/positronic/vendors/gr00t/server.py b/positronic/vendors/gr00t/server.py index b6c579a95..d757c7b51 100644 --- a/positronic/vendors/gr00t/server.py +++ b/positronic/vendors/gr00t/server.py @@ -2,6 +2,7 @@ import os import subprocess from collections.abc import Callable +from enum import Enum, auto from pathlib import Path from typing import Any @@ -27,6 +28,10 @@ logger = logging.getLogger(__name__) +NUMPY_ARRAY = b'nd' +NUMPY_KIND = b'kind' +NUMPY_OBJECT_KIND = b'O' + class MsgSerializer: """N1.7's msgpack-numpy wire format, excluding pickle-bearing object arrays.""" @@ -42,7 +47,9 @@ def from_bytes(data: bytes) -> Any: @staticmethod def decode_custom_classes(obj): if isinstance(obj, dict): - if obj.get(b'nd', obj.get('nd')) and obj.get(b'kind', obj.get('kind')) in (b'O', 'O'): + if obj.get(NUMPY_ARRAY, obj.get(NUMPY_ARRAY.decode())) and obj.get( + NUMPY_KIND, obj.get(NUMPY_KIND.decode()) + ) in (NUMPY_OBJECT_KIND, NUMPY_OBJECT_KIND.decode()): raise ValueError('Object arrays are not supported by the GR00T wire protocol') if obj.get(gr00t.MODALITY_CONFIG): return obj[gr00t.AS_JSON] @@ -55,6 +62,11 @@ def encode_custom_classes(obj): return mnp.encode(obj) +class PingResult(Enum): + SUCCESS = auto() + FAILURE = auto() + + class PolicyClient: """Client for communicating with GR00T N1.7 PolicyServer via ZMQ.""" @@ -72,12 +84,12 @@ def _make_socket(self): socket.connect(f'tcp://{self.host}:{self.port}') return socket - def ping(self) -> bool: + def ping(self) -> PingResult: try: self.call_endpoint(gr00t.PING) - return True + return PingResult.SUCCESS except (zmq.error.ZMQError, RuntimeError): - return False + return PingResult.FAILURE def call_endpoint(self, endpoint: str, data: dict | None = None) -> Any: request: dict = {gr00t.ENDPOINT: endpoint} @@ -142,7 +154,7 @@ def _wait_for_ready(self, on_progress: Callable[[str], None] | None): client = PolicyClient(host='127.0.0.1', port=self.zmq_port, timeout_ms=2000) try: wait_for_subprocess_ready( - client.ping, + lambda: client.ping() is PingResult.SUCCESS, lambda: (self.process.poll() is not None, self.process.returncode), 'gr00t subprocess', on_progress, diff --git a/positronic/vendors/gr00t/tests/test_server.py b/positronic/vendors/gr00t/tests/test_server.py index 3d073abdd..55a325838 100644 --- a/positronic/vendors/gr00t/tests/test_server.py +++ b/positronic/vendors/gr00t/tests/test_server.py @@ -101,8 +101,8 @@ def test_client_can_ping_after_a_transport_failure(monkeypatch, failure): monkeypatch.setattr(gr00t_server.zmq, 'Context', lambda: context) client = gr00t_server.PolicyClient() try: - assert not client.ping() - assert client.ping() + assert client.ping() is gr00t_server.PingResult.FAILURE + assert client.ping() is gr00t_server.PingResult.SUCCESS request = gr00t_server.MsgSerializer.from_bytes(recovered.send.call_args.args[0]) assert request == {'endpoint': 'ping'} finally: From d276ceb3a8301804f5f7e0d02eb14a12386cc127 Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Fri, 11 Sep 2026 13:50:54 +0300 Subject: [PATCH 13/14] Pin the focused GR00T N1.7 fork patch --- docker/Makefile | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docker/Makefile b/docker/Makefile index e40e9ba9c..57bbfa35b 100644 --- a/docker/Makefile +++ b/docker/Makefile @@ -43,7 +43,7 @@ TAG_DREAMZERO_BASE_LATEST := $(IMAGE_NAME_DREAMZERO_BASE):latest # OpenPI base image (pulled from registry) TAG_OPENPI_BASE := $(IMAGE_NAME_OPENPI_BASE):latest -GROOT_REF := e63c6f70257574cd7cf41a5abd24d3f0cab1306c +GROOT_REF := ee9dd54d3f38150253c6c160fb21b18ae9673562 ifndef GROOT_BASE_IMAGE GROOT_BASE_IMAGE := positro/gr00t-base:$(GROOT_REF) build-groot: build-groot-base From c41da6f13b9b1e0215ffb9ac4f59287472b76776 Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Fri, 11 Sep 2026 14:04:34 +0300 Subject: [PATCH 14/14] Pin GR00T fork with Mac cross-build support --- docker/Makefile | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docker/Makefile b/docker/Makefile index 57bbfa35b..94bea2148 100644 --- a/docker/Makefile +++ b/docker/Makefile @@ -43,7 +43,7 @@ TAG_DREAMZERO_BASE_LATEST := $(IMAGE_NAME_DREAMZERO_BASE):latest # OpenPI base image (pulled from registry) TAG_OPENPI_BASE := $(IMAGE_NAME_OPENPI_BASE):latest -GROOT_REF := ee9dd54d3f38150253c6c160fb21b18ae9673562 +GROOT_REF := be79d6244dda302ace1ff7a2cd55239aad2ad109 ifndef GROOT_BASE_IMAGE GROOT_BASE_IMAGE := positro/gr00t-base:$(GROOT_REF) build-groot: build-groot-base