From 1b41639fb8918b0bee677cf0759e1ac2d249e792 Mon Sep 17 00:00:00 2001 From: cc Date: Fri, 4 Sep 2026 12:55:21 -0700 Subject: [PATCH 1/4] feat(imitation): add profile-driven policy backends Generate concrete recorder and rollout ports from typed policy I/O profiles while preserving Blueprint autoconnect. Move LeRobot onto a shared safety runtime and add Dual OpenYAM two-camera collection plus three-camera Amazon ABC-DiT rollout in an isolated locked environment. Tests: 147 affected host tests; 11 isolated LeRobot tests; 3 isolated ABC tests; strict mypy; pre-commit; blueprint registry; wheel and sdist builds. --- .pre-commit-config.yaml | 2 +- MANIFEST.in | 11 + dimos/imitation/cameras.py | 74 ++ dimos/imitation/dataprep/lerobot.py | 4 +- dimos/imitation/policy/abc/module.py | 47 + dimos/imitation/policy/abc/python/VENDORED.md | 11 + .../policy/abc/python/abc_minimal/__init__.py | 1 + .../policy/abc/python/abc_minimal/config.py | 48 + .../policy/abc/python/abc_minimal/dit.py | 895 ++++++++++++++++++ .../abc/python/abc_minimal/fast_inference.py | 118 +++ .../abc/python/abc_minimal/preprocess.py | 71 ++ .../policy/abc/python/dimos_abc/__init__.py | 1 + .../policy/abc/python/dimos_abc/runtime.py | 208 ++++ .../abc/python/dimos_abc/runtime_tests.py | 40 + .../policy/abc/python/pyproject.toml | 56 ++ dimos/imitation/policy/abc/python/uv.lock | 569 +++++++++++ dimos/imitation/policy/abc/test_module.py | 45 + dimos/imitation/policy/backend.py | 53 ++ dimos/imitation/policy/lerobot/README.md | 63 +- .../lerobot/python/dimos_lerobot/runtime.py | 548 ++--------- .../python/dimos_lerobot/runtime_tests.py | 655 +------------ .../policy/lerobot/python/pyproject.toml | 4 + dimos/imitation/policy/lerobot/test_module.py | 69 +- dimos/imitation/policy/module.py | 161 ++++ dimos/imitation/policy/runtime.py | 485 ++++++++++ dimos/imitation/policy/test_runtime.py | 320 +++++++ dimos/imitation/profile.py | 159 ++++ dimos/imitation/test_profile.py | 107 +++ dimos/robot/all_blueprints.py | 1 - .../blueprints/learning_collection.py | 56 ++ .../blueprints/learning_rollout.py | 78 ++ .../dual_openyam/blueprints/teleop.py | 86 +- .../dual_openyam/blueprints/test_learning.py | 86 ++ .../manipulators/dual_openyam/learning.py | 107 +++ .../manipulation/imitation-learning.md | 28 +- pyproject.toml | 6 +- 36 files changed, 4044 insertions(+), 1229 deletions(-) create mode 100644 dimos/imitation/cameras.py create mode 100644 dimos/imitation/policy/abc/module.py create mode 100644 dimos/imitation/policy/abc/python/VENDORED.md create mode 100644 dimos/imitation/policy/abc/python/abc_minimal/__init__.py create mode 100644 dimos/imitation/policy/abc/python/abc_minimal/config.py create mode 100644 dimos/imitation/policy/abc/python/abc_minimal/dit.py create mode 100644 dimos/imitation/policy/abc/python/abc_minimal/fast_inference.py create mode 100644 dimos/imitation/policy/abc/python/abc_minimal/preprocess.py create mode 100644 dimos/imitation/policy/abc/python/dimos_abc/__init__.py create mode 100644 dimos/imitation/policy/abc/python/dimos_abc/runtime.py create mode 100644 dimos/imitation/policy/abc/python/dimos_abc/runtime_tests.py create mode 100644 dimos/imitation/policy/abc/python/pyproject.toml create mode 100644 dimos/imitation/policy/abc/python/uv.lock create mode 100644 dimos/imitation/policy/abc/test_module.py create mode 100644 dimos/imitation/policy/backend.py create mode 100644 dimos/imitation/policy/module.py create mode 100644 dimos/imitation/policy/runtime.py create mode 100644 dimos/imitation/policy/test_runtime.py create mode 100644 dimos/imitation/profile.py create mode 100644 dimos/imitation/test_profile.py create mode 100644 dimos/robot/manipulators/dual_openyam/blueprints/learning_collection.py create mode 100644 dimos/robot/manipulators/dual_openyam/blueprints/learning_rollout.py create mode 100644 dimos/robot/manipulators/dual_openyam/blueprints/test_learning.py create mode 100644 dimos/robot/manipulators/dual_openyam/learning.py diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index de6ab3e28f..2825b53048 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,6 +1,6 @@ default_stages: [pre-commit] default_install_hook_types: [pre-commit, commit-msg] -exclude: (dimos/models/.*)|(deprecated) +exclude: (dimos/models/.*)|(deprecated)|(dimos/imitation/policy/abc/python/abc_minimal/) repos: - repo: https://github.com/Lucas-C/pre-commit-hooks diff --git a/MANIFEST.in b/MANIFEST.in index 912e9ad0a9..e9e13ae147 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -11,6 +11,9 @@ global-exclude .DS_Store recursive-include dimos *.yaml *.yml *.json *.urdf *.html *.css *.js *.svg *.tcss include dimos/imitation/policy/lerobot/python/pyproject.toml include dimos/imitation/policy/lerobot/python/uv.lock +include dimos/imitation/policy/abc/python/pyproject.toml +include dimos/imitation/policy/abc/python/uv.lock +include dimos/imitation/policy/abc/python/VENDORED.md # --- Exclusions (must come after the includes above so they win) --- # Test fixtures must never ship. @@ -45,6 +48,14 @@ prune .pytest_cache prune .ruff_cache prune .vscode prune dimos/web/command-center-extension +prune dimos/imitation/policy/lerobot/python/.mypy_cache +prune dimos/imitation/policy/lerobot/python/.pytest_cache +prune dimos/imitation/policy/lerobot/python/.ruff_cache +prune dimos/imitation/policy/lerobot/python/.venv +prune dimos/imitation/policy/abc/python/.mypy_cache +prune dimos/imitation/policy/abc/python/.pytest_cache +prune dimos/imitation/policy/abc/python/.ruff_cache +prune dimos/imitation/policy/abc/python/.venv global-exclude test_*.py global-exclude conftest.py diff --git a/dimos/imitation/cameras.py b/dimos/imitation/cameras.py new file mode 100644 index 0000000000..153c078d2d --- /dev/null +++ b/dimos/imitation/cameras.py @@ -0,0 +1,74 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Camera blueprints generated from policy image-source declarations.""" + +from __future__ import annotations + +from collections.abc import Mapping + +from dimos.core.coordination.blueprints import Blueprint +from dimos.hardware.sensors.camera.module import CameraModule +from dimos.hardware.sensors.camera.webcam import WebcamConfig +from dimos.imitation.profile import ImageSource, PolicyIOProfile + +CameraDevice = int | str + + +def profile_cameras( + profile: PolicyIOProfile, + devices: Mapping[str, CameraDevice], +) -> tuple[list[Blueprint], list[tuple[str, str, str]]]: + """Build cameras and explicit output remappings for a policy profile.""" + image_sources = { + source.stream: source + for source in profile.observations.values() + if isinstance(source, ImageSource) + } + missing = sorted(set(image_sources) - set(devices)) + unknown = sorted(set(devices) - set(image_sources)) + if missing or unknown: + details = [] + if missing: + details.append(f"missing cameras: {missing}") + if unknown: + details.append(f"unknown cameras: {unknown}") + raise ValueError("; ".join(details)) + + blueprints: list[Blueprint] = [] + remappings: list[tuple[str, str, str]] = [] + for stream_name, source in image_sources.items(): + height, width, _channels = source.shape + instance_name = f"PolicyCamera_{stream_name}" + blueprints.append( + CameraModule.blueprint( + instance_name=instance_name, + hardware=WebcamConfig( + camera_index=devices[stream_name], + width=width, + height=height, + fps=profile.sync.rate_hz, + frame_id_prefix=stream_name, + ), + frame_id=f"{stream_name}_camera_link", + ) + ) + remappings.extend( + [ + (instance_name, "color_image", stream_name), + (instance_name, "camera_info", f"{stream_name}_camera_info"), + (instance_name, "tf", f"{stream_name}_tf"), + ] + ) + return blueprints, remappings diff --git a/dimos/imitation/dataprep/lerobot.py b/dimos/imitation/dataprep/lerobot.py index b6fc404330..0422b709df 100644 --- a/dimos/imitation/dataprep/lerobot.py +++ b/dimos/imitation/dataprep/lerobot.py @@ -33,13 +33,13 @@ Result, ) from dimos.imitation.dataprep.core import DataPrepConfig -from dimos.imitation.policy.lerobot.module import LeRobotPolicyModule +from dimos.imitation.policy.lerobot.module import OpenYamLeRobotPolicy from dimos.utils.cache import cache_usage_guard def lerobot_project() -> Path: """Locate the packaged LeRobot project beside its host contract.""" - source = Path(inspect.getfile(LeRobotPolicyModule)).resolve() + source = Path(inspect.getfile(OpenYamLeRobotPolicy)).resolve() return source.parent / "python" diff --git a/dimos/imitation/policy/abc/module.py b/dimos/imitation/policy/abc/module.py new file mode 100644 index 0000000000..0c58991194 --- /dev/null +++ b/dimos/imitation/policy/abc/module.py @@ -0,0 +1,47 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Dual OpenYAM binding for Amazon's released ABC-DiT checkpoint.""" + +from pathlib import Path + +from pydantic import Field, field_validator + +from dimos.imitation.policy.module import PolicyRolloutConfig, declare_policy_module +from dimos.robot.manipulators.dual_openyam.learning import DUAL_OPENYAM_ABC_IO + + +class AbcPolicyConfig(PolicyRolloutConfig): + """Released ABC-DiT inference settings.""" + + norm_stats_path: str | None = None + diffusion_steps: int = Field(default=10, ge=1) + fast_inference: bool = True + + @field_validator("norm_stats_path") + @classmethod + def resolve_norm_stats_path(cls, value: str | None) -> str | None: + if value is None: + return None + path = Path(value).expanduser() + return str(path.resolve()) if path.exists() else value + + +DualOpenYamAbcPolicy = declare_policy_module( + "DualOpenYamAbcPolicy", + __name__, + DUAL_OPENYAM_ABC_IO, + AbcPolicyConfig, + "dimos_abc.runtime:AbcPolicyRuntime", +) diff --git a/dimos/imitation/policy/abc/python/VENDORED.md b/dimos/imitation/policy/abc/python/VENDORED.md new file mode 100644 index 0000000000..0249a1f13b --- /dev/null +++ b/dimos/imitation/policy/abc/python/VENDORED.md @@ -0,0 +1,11 @@ +# Vendored ABC inference code + +`abc_minimal/dit.py`, `preprocess.py`, and `fast_inference.py` come from +[`amazon-far/abc`](https://github.com/amazon-far/abc) at revision +`6bc6586721cf0c409ccee80f675a28de9b9b2f5e`. `config.py` retains only the two +configuration dataclasses used for inference, and `fast_inference.py` omits the +RTC helper. The upstream project is licensed under Apache-2.0, the same license +as this repository. + +Training, simulation, dataset conversion, visualization, and RTC code are not +vendored. diff --git a/dimos/imitation/policy/abc/python/abc_minimal/__init__.py b/dimos/imitation/policy/abc/python/abc_minimal/__init__.py new file mode 100644 index 0000000000..2bcd1fedfc --- /dev/null +++ b/dimos/imitation/policy/abc/python/abc_minimal/__init__.py @@ -0,0 +1 @@ +"""Vendored ABC-DiT inference subset.""" diff --git a/dimos/imitation/policy/abc/python/abc_minimal/config.py b/dimos/imitation/policy/abc/python/abc_minimal/config.py new file mode 100644 index 0000000000..a31d826d4c --- /dev/null +++ b/dimos/imitation/policy/abc/python/abc_minimal/config.py @@ -0,0 +1,48 @@ +"""Inference configuration extracted from the upstream ABC minimal release.""" + +import os +from dataclasses import dataclass, field +from pathlib import Path + + +REPO_ROOT = Path(__file__).resolve().parents[1] +DEFAULT_CACHE_ROOT = REPO_ROOT / "cache" + + +def default_cache_root() -> Path: + return Path(os.environ.get("ABC_CACHE", str(DEFAULT_CACHE_ROOT))).expanduser() + + +@dataclass +class ClipConfig: + """CLIP ViT-B/32 text asset locations.""" + + cache_dir: str = field(default_factory=lambda: str(Path.home() / ".cache" / "clip")) + model_url: str = ( + "https://openaipublic.azureedge.net/clip/models/" + "40d365715913c9da98579312b702a82c18be219cc2a73407c4526f58eba950af/ViT-B-32.pt" + ) + bpe_url: str = "https://github.com/openai/CLIP/raw/main/clip/bpe_simple_vocab_16e6.txt.gz" + model_name: str = "ViT-B-32.pt" + bpe_name: str = "bpe_simple_vocab_16e6.txt.gz" + + +@dataclass +class DiTConfig: + """ABC-DiT architecture used by the released checkpoint.""" + + hidden_size: int = 1536 + depth: int = 32 + num_heads: int = 24 + mlp_ratio: float = 4.0 + state_dim: int = 14 + action_dim: int = 14 + chunk_length: int = 30 + camera_keys: tuple[str, ...] = ("top", "left", "right") + task_embed_dim: int = 512 + vit_embed_dim: int = 768 + vit_depth: int = 12 + vit_num_heads: int = 12 + vision_pool_num_queries: int = 12 + vision_pool_num_heads: int = 8 + vision_pool_mlp_ratio: int = 4 diff --git a/dimos/imitation/policy/abc/python/abc_minimal/dit.py b/dimos/imitation/policy/abc/python/abc_minimal/dit.py new file mode 100644 index 0000000000..82c03dde00 --- /dev/null +++ b/dimos/imitation/policy/abc/python/abc_minimal/dit.py @@ -0,0 +1,895 @@ +"""ABC-DiT policy implementation for the released bottles-in-bin checkpoints. + +Includes the CLIP text encoder and DINOv3 vision backbone needed to run the model. +""" + +import gzip +import html +import math +import urllib.request +from functools import lru_cache +from pathlib import Path + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F + +from abc_minimal.config import ClipConfig, DiTConfig + +# CLIP ViT-B/32 text encoder. + +SOT_TOKEN = "<|startoftext|>" +EOT_TOKEN = "<|endoftext|>" + + +def task_name_to_prompt(task_name): + """Convert production task names like open_the_pen_caps to CLIP prompt text.""" + return " ".join(task_name.replace("-", " ").replace("_", " ").split()) + + +def _load_clip_text_deps(): + try: + import ftfy + import regex + except ImportError as exc: + raise RuntimeError( + "CLIP text embedding requires the 'ftfy' and 'regex' packages. " + "Run `uv sync` after pulling this version, or install them manually." + ) from exc + return ftfy, regex + + +def _download_if_missing(url, path): + path.parent.mkdir(parents=True, exist_ok=True) + if not path.exists(): + urllib.request.urlretrieve(url, path) + + +def ensure_clip_text_assets(config: ClipConfig): + """Download CLIP ViT-B/32 text assets if needed.""" + root = Path(config.cache_dir).expanduser() + b32_path = root / config.model_name + bpe_path = root / config.bpe_name + _download_if_missing(config.model_url, b32_path) + _download_if_missing(config.bpe_url, bpe_path) + return b32_path, bpe_path + + +@lru_cache() +def _bytes_to_unicode(): + bs = ( + list(range(ord("!"), ord("~") + 1)) + + list(range(ord("¡"), ord("¬") + 1)) + + list(range(ord("®"), ord("ÿ") + 1)) + ) + cs = bs[:] + n = 0 + for b in range(256): + if b not in bs: + bs.append(b) + cs.append(256 + n) + n += 1 + return dict(zip(bs, [chr(c) for c in cs])) + + +def _get_pairs(word): + return set(zip(word, word[1:])) + + +class CLIPBPETokenizer: + """Small copy of OpenAI CLIP's BPE tokenizer, scoped to text encoding.""" + + def __init__(self, bpe_path): + _ftfy, regex = _load_clip_text_deps() + self.ftfy = _ftfy + self.regex = regex + with gzip.open(bpe_path, "rt", encoding="utf-8") as f: + merges = [tuple(l.split()) for l in f.read().split("\n")[1 : 49152 - 256 - 2 + 1]] + self.byte_encoder = _bytes_to_unicode() + vocab = list(self.byte_encoder.values()) + vocab = vocab + [v + "" for v in vocab] + for merge in merges: + vocab.append("".join(merge)) + vocab.extend([SOT_TOKEN, EOT_TOKEN]) + self.encoder = {v: i for i, v in enumerate(vocab)} + self.bpe_ranks = dict(zip(merges, range(len(merges)))) + self.cache = {SOT_TOKEN: SOT_TOKEN, EOT_TOKEN: EOT_TOKEN} + self.pat = regex.compile( + r"<\|startoftext\|>|<\|endoftext\|>|\'s|\'t|\'re|\'ve|\'m|\'ll|\'d|" + r"[\p{L}]+|[\p{N}]|[^\s\p{L}\p{N}]+", + regex.IGNORECASE, + ) + + def _basic_clean(self, text): + return html.unescape(html.unescape(self.ftfy.fix_text(text))).strip() + + def _whitespace_clean(self, text): + return self.regex.sub(r"\s+", " ", text).strip() + + def bpe(self, token): + if token in self.cache: + return self.cache[token] + word = tuple(token[:-1]) + (token[-1] + "",) + pairs = _get_pairs(word) + if not pairs: + return token + "" + while True: + bigram = min(pairs, key=lambda p: self.bpe_ranks.get(p, float("inf"))) + if bigram not in self.bpe_ranks: + break + first, second = bigram + new_word = [] + i = 0 + while i < len(word): + try: + j = word.index(first, i) + except ValueError: + new_word.extend(word[i:]) + break + new_word.extend(word[i:j]) + if word[j] == first and j < len(word) - 1 and word[j + 1] == second: + new_word.append(first + second) + i = j + 2 + else: + new_word.append(word[j]) + i = j + 1 + word = tuple(new_word) + if len(word) == 1: + break + pairs = _get_pairs(word) + out = " ".join(word) + self.cache[token] = out + return out + + def encode(self, text): + text = self._whitespace_clean(self._basic_clean(text)).lower() + tokens = [] + for token in self.regex.findall(self.pat, text): + token = "".join(self.byte_encoder[b] for b in token.encode("utf-8")) + tokens.extend(self.encoder[piece] for piece in self.bpe(token).split(" ")) + return tokens + + +class CLIPQuickGELU(nn.Module): + def forward(self, x): + return x * torch.sigmoid(1.702 * x) + + +class CLIPTextBlock(nn.Module): + def __init__(self, width, heads, mask): + super().__init__() + self.attn = nn.MultiheadAttention(width, heads) + self.ln_1 = nn.LayerNorm(width) + self.mlp = nn.Sequential() + self.mlp.add_module("c_fc", nn.Linear(width, width * 4)) + self.mlp.add_module("gelu", CLIPQuickGELU()) + self.mlp.add_module("c_proj", nn.Linear(width * 4, width)) + self.ln_2 = nn.LayerNorm(width) + self.register_buffer("mask", mask, persistent=False) + + def forward(self, x): + x_ln = self.ln_1(x) + x = x + self.attn(x_ln, x_ln, x_ln, need_weights=False, attn_mask=self.mask)[0] + x = x + self.mlp(self.ln_2(x)) + return x + + +class CLIPTextTower(nn.Module): + def __init__(self, state_dict): + super().__init__() + embed_dim = state_dict["text_projection"].shape[1] + context_length = state_dict["positional_embedding"].shape[0] + vocab_size = state_dict["token_embedding.weight"].shape[0] + width = state_dict["ln_final.weight"].shape[0] + heads = width // 64 + layers = len( + {k.split(".")[2] for k in state_dict if k.startswith("transformer.resblocks")} + ) + mask = torch.empty(context_length, context_length).fill_(float("-inf")).triu_(1) + + self.context_length = context_length + self.token_embedding = nn.Embedding(vocab_size, width) + self.positional_embedding = nn.Parameter(torch.empty(context_length, width)) + self.transformer = nn.Module() + self.transformer.resblocks = nn.Sequential( + *[CLIPTextBlock(width, heads, mask) for _ in range(layers)] + ) + self.ln_final = nn.LayerNorm(width) + self.text_projection = nn.Parameter(torch.empty(width, embed_dim)) + + def forward(self, text): + x = self.token_embedding(text) + self.positional_embedding + x = x.permute(1, 0, 2) + x = self.transformer.resblocks(x) + x = x.permute(1, 0, 2) + x = self.ln_final(x) + return x[torch.arange(x.shape[0], device=x.device), text.argmax(dim=-1)] @ self.text_projection + + +class CLIPTextEmbedder: + """OpenAI CLIP ViT-B/32 text encoder that returns normalized 512-d vectors. + Holds a CPU memo cache keyed by prompt so repeats skip BPE+transformer.""" + + def __init__(self, config: ClipConfig, device="cpu"): + b32_path, bpe_path = ensure_clip_text_assets(config) + self.device = torch.device(device) + try: + state_dict = torch.jit.load(str(b32_path), map_location="cpu").state_dict() + except RuntimeError: + state_dict = torch.load(b32_path, map_location="cpu", weights_only=False) + self.tokenizer = CLIPBPETokenizer(bpe_path) + self.model = CLIPTextTower(state_dict).eval().to(self.device) + text_keys = { + k: v + for k, v in state_dict.items() + if k.startswith( + ( + "token_embedding", + "positional_embedding", + "transformer.resblocks", + "ln_final", + "text_projection", + ) + ) + } + missing, unexpected = self.model.load_state_dict(text_keys, strict=False) + if missing or unexpected: + raise RuntimeError(f"bad CLIP text weights: missing={missing} unexpected={unexpected}") + self._cache = {} + + @torch.no_grad() + def encode(self, texts): + if isinstance(texts, str): + texts = [texts] + fresh = [t for t in dict.fromkeys(texts) if t not in self._cache] + if fresh: + context = torch.zeros( + len(fresh), self.model.context_length, dtype=torch.long, device=self.device + ) + for i, text in enumerate(fresh): + token_ids = [ + self.tokenizer.encoder[SOT_TOKEN], + *self.tokenizer.encode(text), + self.tokenizer.encoder[EOT_TOKEN], + ] + if len(token_ids) > self.model.context_length: + raise RuntimeError( + f"Input {text!r} is too long for CLIP context length " + f"{self.model.context_length}" + ) + context[i, : len(token_ids)] = torch.tensor( + token_ids, dtype=torch.long, device=self.device + ) + features = self.model(context) + features = features / features.norm(dim=-1, keepdim=True) + for i, text in enumerate(fresh): + self._cache[text] = features[i].cpu() + out = torch.stack([self._cache[t] for t in texts], dim=0) + return out.to(self.device) + + +def encode_clip_text(texts, config: ClipConfig, device="cpu"): + """Encode exact prompt text with OpenAI CLIP ViT-B/32.""" + return CLIPTextEmbedder(config, device=device).encode(texts) + + +def encode_clip_task_name(task_names, config: ClipConfig, device="cpu"): + """Encode task names after production-style dash/underscore replacement.""" + if isinstance(task_names, str): + task_names = [task_names] + return encode_clip_text([task_name_to_prompt(t) for t in task_names], config, device=device) + + +# DINOv3 ViT-B/16 vision encoder. + + +def _rope_rotate_half(x): + x1, x2 = x.chunk(2, dim=-1) + return torch.cat([-x2, x1], dim=-1) + + +class DinoRope(nn.Module): + """RoPE over the 2D patch grid (base=100, separate coord normalization). + + rescale_coords=2 applies a random log-uniform rescale of the coordinates + during training only — part of the pretraining distribution, kept for + finetuning fidelity. + """ + + def __init__(self, embed_dim, num_heads, base=100.0, rescale_coords=2.0): + super().__init__() + d_head = embed_dim // num_heads + self.d_head = d_head + self.rescale_coords = rescale_coords + self.register_buffer("periods", torch.empty(d_head // 4), persistent=True) + with torch.no_grad(): + self.periods.copy_( + base ** (2 * torch.arange(d_head // 4, dtype=torch.float32) / (d_head // 2)) + ) + + def forward(self, H, W): + dev = self.periods.device + coords_h = torch.arange(0.5, H, device=dev, dtype=torch.float32) / H + coords_w = torch.arange(0.5, W, device=dev, dtype=torch.float32) / W + coords = torch.stack(torch.meshgrid(coords_h, coords_w, indexing="ij"), dim=-1) + coords = coords.flatten(0, 1) + coords = 2.0 * coords - 1.0 + if self.training and self.rescale_coords is not None: + r = np.log(self.rescale_coords) + rescale = torch.empty(1, device=dev).uniform_(-r, r).exp() + coords = coords * rescale + angles = 2 * math.pi * coords[:, :, None] / self.periods[None, None, :] + angles = angles.flatten(1, 2).tile(2) + return torch.sin(angles), torch.cos(angles) + + +class LinearKMaskedBias(nn.Linear): + """qkv Linear whose k-third of the bias is masked to zero (DINOv3 quirk).""" + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.register_buffer("bias_mask", torch.full_like(self.bias, math.nan)) + + def forward(self, x): + return F.linear(x, self.weight, self.bias * self.bias_mask.to(self.bias.dtype)) + + +class DinoAttention(nn.Module): + def __init__(self, dim, num_heads): + super().__init__() + self.num_heads = num_heads + self.qkv = LinearKMaskedBias(dim, dim * 3, bias=True) + self.proj = nn.Linear(dim, dim, bias=True) + + def forward(self, x, rope=None): + B, N, C = x.shape + qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads) + q, k, v = [t.transpose(1, 2) for t in torch.unbind(qkv, 2)] + if rope is not None: + sin, cos = rope + n_prefix = N - sin.shape[-2] # cls + storage tokens are not rotated + q_dt, k_dt = q.dtype, k.dtype + q, k = q.to(sin.dtype), k.to(sin.dtype) + q = torch.cat( + [q[:, :, :n_prefix], q[:, :, n_prefix:] * cos + _rope_rotate_half(q[:, :, n_prefix:]) * sin], + dim=-2, + ) + k = torch.cat( + [k[:, :, :n_prefix], k[:, :, n_prefix:] * cos + _rope_rotate_half(k[:, :, n_prefix:]) * sin], + dim=-2, + ) + q, k = q.to(q_dt), k.to(k_dt) + x = F.scaled_dot_product_attention(q, k, v) + return self.proj(x.transpose(1, 2).reshape(B, N, C)) + + +class LayerScale(nn.Module): + def __init__(self, dim, init_values=1e-5): + super().__init__() + self.gamma = nn.Parameter(init_values * torch.ones(dim)) + + def forward(self, x): + return x * self.gamma + + +class DinoMlp(nn.Module): + def __init__(self, dim, hidden): + super().__init__() + self.fc1 = nn.Linear(dim, hidden) + self.act = nn.GELU() + self.fc2 = nn.Linear(hidden, dim) + + def forward(self, x): + return self.fc2(self.act(self.fc1(x))) + + +class DinoBlock(nn.Module): + def __init__(self, dim, num_heads, ffn_ratio=4.0): + super().__init__() + self.norm1 = nn.LayerNorm(dim, eps=1e-5) + self.attn = DinoAttention(dim, num_heads) + self.ls1 = LayerScale(dim) + self.norm2 = nn.LayerNorm(dim, eps=1e-5) + self.mlp = DinoMlp(dim, int(dim * ffn_ratio)) + self.ls2 = LayerScale(dim) + + def forward(self, x, rope=None): + x = x + self.ls1(self.attn(self.norm1(x), rope=rope)) + x = x + self.ls2(self.mlp(self.norm2(x))) + return x + + +class DinoPatchEmbed(nn.Module): + def __init__(self, patch_size=16, in_chans=3, embed_dim=768): + super().__init__() + self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size) + + def forward(self, x): + x = self.proj(x) # (B, D, H/16, W/16) + return x.flatten(2).transpose(1, 2), x.shape[2], x.shape[3] + + +class DinoVisionTransformer(nn.Module): + """DINOv3 ViT-B/16 with 4 storage tokens. encode_image_tokens() returns + (B, 1+196, 768) = CLS + patch tokens (storage tokens dropped), matching + the production vision backbone interface.""" + + N_STORAGE_TOKENS = 4 + + def __init__(self, embed_dim, depth, num_heads): + super().__init__() + self.embed_dim = embed_dim + self.patch_embed = DinoPatchEmbed(embed_dim=embed_dim) + self.cls_token = nn.Parameter(torch.empty(1, 1, embed_dim)) + self.storage_tokens = nn.Parameter(torch.empty(1, self.N_STORAGE_TOKENS, embed_dim)) + self.mask_token = nn.Parameter(torch.empty(1, embed_dim)) + self.rope_embed = DinoRope(embed_dim, num_heads) + self.blocks = nn.ModuleList(DinoBlock(embed_dim, num_heads) for _ in range(depth)) + self.norm = nn.LayerNorm(embed_dim, eps=1e-5) + self.init_weights() + + def init_weights(self): + """Match production's models/dinov3/vision_transformer.py:init_weights_vit. + Crucially, this fills `bias_mask` (otherwise NaN-initialized) so that + the K-third of every qkv bias is masked to 0 — without this, a fresh + DINOv3 produces NaN on its very first forward pass.""" + nn.init.normal_(self.cls_token, std=0.02) + nn.init.normal_(self.storage_tokens, std=0.02) + nn.init.zeros_(self.mask_token) + for m in self.modules(): + if isinstance(m, nn.Linear): + nn.init.trunc_normal_(m.weight, std=0.02) + if m.bias is not None: + nn.init.zeros_(m.bias) + if isinstance(m, LinearKMaskedBias): + o = m.out_features + m.bias_mask.fill_(1) + m.bias_mask[o // 3 : 2 * o // 3].fill_(0) + elif isinstance(m, nn.LayerNorm): + m.reset_parameters() + elif isinstance(m, LayerScale): + nn.init.constant_(m.gamma, 1e-5) + elif isinstance(m, DinoPatchEmbed): + # Match nn.Conv2d default + m.proj.reset_parameters() + + def encode_image_tokens(self, images): + x, H, W = self.patch_embed(images) + B = x.shape[0] + cls_token = self.cls_token + 0 * self.mask_token # production quirk, kept + x = torch.cat( + [cls_token.expand(B, -1, -1), self.storage_tokens.expand(B, -1, -1), x], dim=1 + ) + rope = self.rope_embed(H, W) + for blk in self.blocks: + x = blk(x, rope=rope) + x = self.norm(x) + cls_out = x[:, :1] + patches = x[:, 1 + self.N_STORAGE_TOKENS :] + return torch.cat([cls_out, patches], dim=1) + + +class DinoVisionBackbone(nn.Module): + """Wrapper around DinoVisionTransformer with an optional bf16-autocast + forward path. The wrapper keeps the production checkpoint key layout + (`img_backbone.dinov3_model.*`) so the slim 200k checkpoint loads with + zero missing/unexpected keys. Set bf16 with set_bfloat16(True): the + DINO forward then runs under autocast(bf16) on CUDA, cutting + vision-encoder activation memory roughly in half. Tokens are cast back + to fp32 on the way out so the surrounding DiT stays dtype-stable. + """ + + def __init__(self, config: DiTConfig): + super().__init__() + self.dinov3_model = DinoVisionTransformer( + embed_dim=config.vit_embed_dim, + depth=config.vit_depth, + num_heads=config.vit_num_heads, + ) + self.bfloat16 = False + + def set_bfloat16(self, enabled: bool = True): + self.bfloat16 = bool(enabled) + + def encode_image_tokens(self, images): + if self.bfloat16 and images.is_cuda: + with torch.autocast(device_type="cuda", dtype=torch.bfloat16): + tokens = self.dinov3_model.encode_image_tokens(images) + return tokens.to(torch.float32) + return self.dinov3_model.encode_image_tokens(images) + + +# ABC-DiT policy. + + +def modulate(x, shift, scale): + if shift.ndim == 2: + shift = shift.unsqueeze(1) + scale = scale.unsqueeze(1) + return x * (1 + scale) + shift + + +def gate_residual(gate, residual): + if gate.ndim == 2: + gate = gate.unsqueeze(1) + return gate * residual + + +def get_1d_sincos_pos_embed(embed_dim, length): + omega = np.arange(embed_dim // 2, dtype=np.float64) + omega /= embed_dim / 2.0 + omega = 1.0 / 10000**omega + out = np.einsum("m,d->md", np.arange(length, dtype=np.float64), omega) + return np.concatenate([np.sin(out), np.cos(out)], axis=1) + + +class TimestepEmbedder(nn.Module): + def __init__(self, hidden_size, frequency_embedding_size=256): + super().__init__() + self.mlp = nn.Sequential( + nn.Linear(frequency_embedding_size, hidden_size, bias=True), + nn.SiLU(), + nn.Linear(hidden_size, hidden_size, bias=True), + ) + self.frequency_embedding_size = frequency_embedding_size + half = frequency_embedding_size // 2 + freqs = torch.exp( + -math.log(10000) * torch.arange(half, dtype=torch.float32) / half + ) + self.register_buffer("freqs", freqs, persistent=False) + + def timestep_embedding(self, t): + freqs = self.freqs + if freqs.device != t.device: + freqs = freqs.to(device=t.device) + args = t[:, None].float() * freqs[None] + return torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + + def forward(self, t): + t_shape = t.shape + t_freq = self.timestep_embedding(t.reshape(-1)) + t_emb = self.mlp(t_freq.to(self.mlp[0].weight.dtype)) + return t_emb.reshape(*t_shape, -1) + + +class DiTAttention(nn.Module): + """Self-attention over action tokens (timm-equivalent, qkv_bias=True).""" + + def __init__(self, dim, num_heads): + super().__init__() + self.num_heads = num_heads + self.qkv = nn.Linear(dim, dim * 3, bias=True) + self.proj = nn.Linear(dim, dim, bias=True) + + def forward(self, x): + B, N, C = x.shape + qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads) + q, k, v = qkv.permute(2, 0, 3, 1, 4).unbind(0) + x = F.scaled_dot_product_attention(q, k, v) + return self.proj(x.transpose(1, 2).reshape(B, N, C)) + + +class DiTMlp(nn.Module): + def __init__(self, dim, hidden): + super().__init__() + self.fc1 = nn.Linear(dim, hidden) + self.act = nn.GELU(approximate="tanh") + self.fc2 = nn.Linear(hidden, dim) + + def forward(self, x): + return self.fc2(self.act(self.fc1(x))) + + +class DiTBlock(nn.Module): + """AdaLN-Zero DiT block with vision cross-attention (9-way modulation).""" + + def __init__(self, hidden_size, num_heads, mlp_ratio=4.0): + super().__init__() + self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.attn = DiTAttention(hidden_size, num_heads) + self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.mlp = DiTMlp(hidden_size, int(hidden_size * mlp_ratio)) + self.norm_xattn = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.norm_xattn_kv = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.cross_attn = nn.MultiheadAttention(hidden_size, num_heads, batch_first=True) + self.adaLN_modulation = nn.Sequential( + nn.SiLU(), nn.Linear(hidden_size, 9 * hidden_size, bias=True) + ) + + def forward(self, x, c, vision_tokens): + ( + shift_msa, scale_msa, gate_msa, + shift_xattn, scale_xattn, gate_xattn, + shift_mlp, scale_mlp, gate_mlp, + ) = self.adaLN_modulation(c).chunk(9, dim=-1) + + x = x + gate_residual(gate_msa, self.attn(modulate(self.norm1(x), shift_msa, scale_msa))) + + x_normed = modulate(self.norm_xattn(x), shift_xattn, scale_xattn) + kv = self.norm_xattn_kv(vision_tokens) + xattn_out, _ = self.cross_attn(x_normed, kv, kv, need_weights=False) + x = x + gate_residual(gate_xattn, xattn_out) + + x = x + gate_residual(gate_mlp, self.mlp(modulate(self.norm2(x), shift_mlp, scale_mlp))) + return x + + +class FinalLayer(nn.Module): + def __init__(self, hidden_size, action_dim): + super().__init__() + self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.linear = nn.Linear(hidden_size, action_dim, bias=True) + self.adaLN_modulation = nn.Sequential( + nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size, bias=True) + ) + + def forward(self, x, c): + shift, scale = self.adaLN_modulation(c).chunk(2, dim=-1) + return self.linear(modulate(self.norm_final(x), shift, scale)) + + +class PoolMlp(nn.Module): + def __init__(self, in_dim, hidden_dim): + super().__init__() + self.fc1 = nn.Linear(in_dim, hidden_dim) + self.act = nn.GELU() + self.fc2 = nn.Linear(hidden_dim, in_dim) + + def forward(self, x): + return self.fc2(self.act(self.fc1(x))) + + +class AttentionPoolBlock(nn.Module): + """Learnable queries cross-attend to ViT tokens (per camera).""" + + def __init__(self, embed_dim, num_heads, mlp_ratio=4): + super().__init__() + self.ln_1 = nn.LayerNorm(embed_dim) + self.attention = nn.MultiheadAttention(embed_dim, num_heads, batch_first=True) + self.ln_2 = nn.LayerNorm(embed_dim) + self.mlp = PoolMlp(embed_dim, int(mlp_ratio * embed_dim)) + + def forward(self, x, queries): + x_kv = self.ln_1(x) + x_q = self.ln_1(queries) + out, _ = self.attention(x_q, x_kv, x_kv, need_weights=False) + return self.mlp(self.ln_2(out)) + out + + +class DiTPolicy(nn.Module): + """Minimal ABC-DiT: loads the production dit_xL pretraining checkpoint.""" + + def __init__(self, config: DiTConfig): + super().__init__() + self.config = config + H = config.hidden_size + self.camera_keys = list(config.camera_keys) + self.chunk_length = config.chunk_length + self.action_dim = config.action_dim + + self.x_embedder = nn.Linear(config.state_dim, H) + self.y_embedder = nn.Linear(config.action_dim, H) + # Checkpoint compatibility only; unused in forward. + self.img_proj = nn.Linear(config.vit_embed_dim, H) + self.img_proj.requires_grad_(False) + self.t_embedder = TimestepEmbedder(H) + self.pos_embed = nn.Parameter(torch.zeros(1, config.chunk_length, H), requires_grad=False) + + self.img_backbone = DinoVisionBackbone(config) + + self.apool_queries = nn.ParameterDict( + { + cam: nn.Parameter( + torch.randn(1, config.vision_pool_num_queries, config.vit_embed_dim) * 0.02 + ) + for cam in self.camera_keys + } + ) + self.apool = nn.ModuleDict( + { + cam: AttentionPoolBlock( + config.vit_embed_dim, + config.vision_pool_num_heads, + config.vision_pool_mlp_ratio, + ) + for cam in self.camera_keys + } + ) + self.vision_tokens_proj = nn.Linear(config.vit_embed_dim, H) + self.vision_camera_embed = nn.Embedding(len(self.camera_keys), H) + + self.task_to_hidden = nn.Linear(config.task_embed_dim, H) + self.blocks = nn.ModuleList( + DiTBlock(H, config.num_heads, config.mlp_ratio) for _ in range(config.depth) + ) + self.final_layer = FinalLayer(H, config.action_dim) + + # cond = [state, task, timestep] -> hidden (vision goes via cross-attn) + self.cond_proj = nn.Sequential( + nn.Linear(3 * H, H), nn.SiLU(), nn.Linear(H, H), nn.LayerNorm(H) + ) + + self.register_buffer("clip_mean", torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1)) + self.register_buffer("clip_std", torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1)) + + pos = get_1d_sincos_pos_embed(H, config.chunk_length) + self.pos_embed.data.copy_(torch.from_numpy(pos).float().unsqueeze(0)) + + def build_vision_tokens(self, images): + """images: dict cam -> (B, 3, 224, 224), already ImageNet-normalized. + Returns (B, num_cameras * queries, hidden).""" + pooled = [] + for cam in self.camera_keys: + tokens = self.img_backbone.encode_image_tokens(images[cam]) + tokens = tokens.to(self.apool_queries[cam].dtype) + queries = self.apool_queries[cam].expand(tokens.shape[0], -1, -1) + pooled.append(self.apool[cam](tokens, queries)) + tokens_by_camera = torch.stack(pooled, dim=1) # (B, Nc, K, vit_dim) + B, Nc, K, D = tokens_by_camera.shape + vision_tokens = self.vision_tokens_proj( + tokens_by_camera.reshape(B * Nc * K, D).to(self.vision_tokens_proj.weight.dtype) + ).reshape(B, Nc, K, -1) + cam_emb = self.vision_camera_embed(torch.arange(Nc, device=vision_tokens.device)) + vision_tokens = vision_tokens + cam_emb[None, :, None, :] + return vision_tokens.reshape(B, Nc * K, -1) + + def compute_cond(self, state, task_vec_clip, t_cond): + """state (B,14); task_vec_clip (B,512); t_cond (B,) or (B,T). + Returns conditioning c: (B,H) or (B,T,H).""" + model_dtype = self.x_embedder.weight.dtype + cond_dtype = self.cond_proj[0].weight.dtype + st_vec = self.x_embedder(state.to(model_dtype)) + task_vec_h = self.task_to_hidden(task_vec_clip.to(self.task_to_hidden.weight.dtype)) + task_vec_h = task_vec_h.to(model_dtype) + t_vec = self.t_embedder(t_cond.to(model_dtype)) + cond_parts = [st_vec, task_vec_h, t_vec] + if t_vec.ndim == 3: + T = t_vec.shape[1] + cond_parts = [ + p.unsqueeze(1).expand(-1, T, -1) if p.ndim == 2 else p for p in cond_parts + ] + cond_concat = torch.cat(cond_parts, dim=-1).to(cond_dtype) + if cond_dtype == torch.float32 and cond_concat.is_cuda: + with torch.autocast(device_type="cuda", enabled=False): + return self.cond_proj(cond_concat).to(model_dtype) + return self.cond_proj(cond_concat).to(model_dtype) + + def predict_velocity(self, x_t, c, vision_tokens): + z = self.y_embedder(x_t) + self.pos_embed.data[:, : x_t.shape[1], :] + for block in self.blocks: + z = block(z, c, vision_tokens) + return self.final_layer(z, c) + + def forward( + self, + batch, + noise=None, + t=None, + max_action_prefix=0, + prefix_conditioning_prob=1.0, + prefix_noise_scale=0.0, + ): + """Flow-matching training loss with optional action-prefix conditioning. + batch: state (B,14), images dict, actions (B,30,14), task_vec_clip (B,512), + optional state_is_masked (B,) bool.""" + state = batch["state"] + actions = batch["actions"] + N, T_chunk, D_action = actions.shape + + if noise is None: + noise = torch.randn_like(actions) + if t is None: + t = torch.rand(N, 1, 1, device=state.device, dtype=actions.dtype) + + if max_action_prefix > 0: + apply_prefix = torch.rand(N, device=state.device) < prefix_conditioning_prob + if "state_is_masked" in batch: + apply_prefix = apply_prefix & ~batch["state_is_masked"].to(state.device) + delay = torch.randint(0, max_action_prefix, (N,), device=state.device) + delay = torch.where(apply_prefix, delay, torch.zeros_like(delay)) + prefix_mask = torch.arange(T_chunk, device=state.device)[None, :] < delay[:, None] + prefix_mask_expanded = prefix_mask.unsqueeze(-1) + t_per_pos = torch.where(prefix_mask_expanded, torch.zeros_like(t), t) + else: + prefix_mask_expanded = None + t_per_pos = t + + x_t = (1 - t_per_pos) * actions + t_per_pos * noise + if prefix_noise_scale > 0.0 and prefix_mask_expanded is not None: + x_t = x_t + prefix_mask_expanded.float() * torch.randn_like(x_t) * prefix_noise_scale + + vision_tokens = self.build_vision_tokens(batch["images"]) + t_cond = t_per_pos.squeeze(-1) if prefix_mask_expanded is not None else t[:, 0, 0] + c = self.compute_cond(state, batch["task_vec_clip"], t_cond) + v_t = self.predict_velocity(x_t, c, vision_tokens) + + u_t = noise - actions + if prefix_mask_expanded is not None: + postfix_mask = ~prefix_mask_expanded + masked_loss = ((u_t - v_t) ** 2) * postfix_mask.float() + return masked_loss.sum() / (postfix_mask.float().sum() * D_action + 1e-8) + return F.mse_loss(u_t, v_t) + + @torch.no_grad() + def sample_actions(self, batch, num_steps=10, noise=None): + """Euler flow integration from noise to actions (production tau=1 path). + Vision tokens and the static conditioning are computed once and reused + across steps, like production infer().""" + state = batch["state"] + B = state.shape[0] + model_dtype = self.y_embedder.weight.dtype + if noise is None: + noise = torch.randn( + B, + self.chunk_length, + self.action_dim, + device=state.device, + dtype=model_dtype, + ) + x_t = noise.to(device=state.device, dtype=model_dtype) + vision_tokens = self.build_vision_tokens(batch["images"]) + dt = -1.0 / num_steps + for i in range(num_steps): + t = torch.full((B,), 1.0 + i * dt, device=state.device, dtype=model_dtype) + c = self.compute_cond(state, batch["task_vec_clip"], t) + v = self.predict_velocity(x_t, c, vision_tokens) + x_t = x_t + v * dt + return x_t + + @torch.no_grad() + def sample_actions_rtc(self, batch, action_prefix, prefix_length: int, num_steps=10, noise=None): + """Euler sampling with per-position action-prefix conditioning.""" + state = batch["state"] + B = state.shape[0] + model_dtype = self.y_embedder.weight.dtype + if noise is None: + noise = torch.randn( + B, + self.chunk_length, + self.action_dim, + device=state.device, + dtype=model_dtype, + ) + x_t = noise.to(device=state.device, dtype=model_dtype) + action_prefix = action_prefix.to(device=state.device, dtype=model_dtype) + + prefix_pos = torch.arange(self.chunk_length, device=state.device) < prefix_length + prefix_mask = prefix_pos.view(1, self.chunk_length, 1).expand_as(x_t) + prefix_t_mask = prefix_pos.view(1, self.chunk_length).expand(B, self.chunk_length) + x_t = torch.where(prefix_mask, action_prefix, x_t) + + vision_tokens = self.build_vision_tokens(batch["images"]) + dt = -1.0 / num_steps + for i in range(num_steps): + t = torch.full( + (B, self.chunk_length), + 1.0 + i * dt, + device=state.device, + dtype=model_dtype, + ) + t = torch.where(prefix_t_mask, torch.zeros_like(t), t) + c = self.compute_cond(state, batch["task_vec_clip"], t) + v = self.predict_velocity(x_t, c, vision_tokens) + x_t = x_t + v * dt + x_t = torch.where(prefix_mask, action_prefix, x_t) + return x_t + + +def load_pretrained(model, ckpt_path): + """Load the slim production checkpoint (model-only, prefixes stripped).""" + ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False, mmap=True) + sd = ckpt["model"] if "model" in ckpt else ckpt + sd = {k[len("_orig_mod."):] if k.startswith("_orig_mod.") else k: v for k, v in sd.items()} + missing, unexpected = model.load_state_dict(sd, strict=False) + if unexpected: + raise RuntimeError(f"unexpected checkpoint keys: {unexpected[:8]}") + if missing: + raise RuntimeError(f"missing checkpoint keys: {missing[:8]}") + return ckpt + + +if __name__ == "__main__": + model = DiTPolicy(DiTConfig()) + n_params = sum(p.numel() for p in model.parameters()) + print(f"DiTPolicy built: {n_params / 1e9:.3f}B params") diff --git a/dimos/imitation/policy/abc/python/abc_minimal/fast_inference.py b/dimos/imitation/policy/abc/python/abc_minimal/fast_inference.py new file mode 100644 index 0000000000..42031c8cc3 --- /dev/null +++ b/dimos/imitation/policy/abc/python/abc_minimal/fast_inference.py @@ -0,0 +1,118 @@ +"""CUDA graph helpers for fast ABC-DiT policy inference.""" + +from __future__ import annotations + +from typing import Any, Protocol + +import numpy as np +import torch + +from abc_minimal.preprocess import normalize, resize_pad_normalize, unnormalize + + +class _PolicyForFastInference(Protocol): + model: Any + device: torch.device + config: Any + task_vec: torch.Tensor + norm_stats: dict[str, Any] + diffusion_steps: int + +class FastInferenceGraph: + """Fixed-shape outer CUDA graph over `model.sample_actions`. + + Captures one bf16 batch on GPU; subsequent .infer() calls just memcpy + fresh inputs into the static tensors and replay the graph. + """ + + def __init__(self, policy: _PolicyForFastInference): + self.policy = policy + self.model = policy.model + self.device = policy.device + self.dtype = torch.bfloat16 + self.graph = torch.cuda.CUDAGraph() + self.output: torch.Tensor | None = None + m = policy.config.model + + self.static_state = torch.empty(1, m.state_dim, device=self.device, dtype=self.dtype) + self.static_noise = torch.empty( + 1, m.chunk_length, m.action_dim, device=self.device, dtype=self.dtype + ) + self.static_images = { + cam: torch.empty(1, 3, 224, 224, device=self.device, dtype=self.dtype) + for cam in m.camera_keys + } + self.static_task_vec = policy.task_vec.to(device=self.device, dtype=self.dtype).clone() + self.batch = { + "state": self.static_state, + "actions": torch.zeros( + 1, m.chunk_length, m.action_dim, device=self.device, dtype=self.dtype + ), + "images": self.static_images, + "task_vec_clip": self.static_task_vec, + } + + def _copy_inputs(self, obs: dict[str, Any], noise: np.ndarray | None) -> None: + m = self.policy.config.model + state = normalize( + np.asarray(obs["state"], dtype=np.float32), self.policy.norm_stats["state"] + ) + self.static_state.copy_( + torch.from_numpy(state[None]).to(device=self.device, dtype=self.dtype) + ) + if noise is None: + self.static_noise.normal_() + else: + noise_arr = noise[None].astype(np.float32, copy=False) + if noise_arr.shape != (1, m.chunk_length, m.action_dim): + raise ValueError( + f"fast inference expects noise shape " + f"{(m.chunk_length, m.action_dim)}, got {noise.shape}" + ) + self.static_noise.copy_( + torch.from_numpy(noise_arr).to(device=self.device, dtype=self.dtype) + ) + for cam in m.camera_keys: + self.static_images[cam].copy_( + resize_pad_normalize(obs["images"][cam]) + .unsqueeze(0) + .to(device=self.device, dtype=self.dtype) + ) + + def capture( + self, + warmup_obs: dict[str, Any], + warmup_noise: np.ndarray | None, + replay_warmups: int, + ) -> None: + self._copy_inputs(warmup_obs, warmup_noise) + + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(5): + self.output = self.model.sample_actions( + self.batch, + num_steps=self.policy.diffusion_steps, + noise=self.static_noise, + ) + torch.cuda.current_stream().wait_stream(stream) + + with torch.cuda.graph(self.graph): + self.output = self.model.sample_actions( + self.batch, num_steps=self.policy.diffusion_steps, noise=self.static_noise + ) + + for _ in range(replay_warmups): + self._copy_inputs(warmup_obs, warmup_noise) + self.graph.replay() + assert self.output is not None + _ = self.output[0].float().detach().cpu().numpy() + torch.cuda.synchronize() + + def infer(self, obs: dict[str, Any], noise: np.ndarray | None) -> np.ndarray: + self._copy_inputs(obs, noise) + self.graph.replay() + assert self.output is not None + actions_np = self.output[0].float().detach().cpu().numpy() + return unnormalize(actions_np, self.policy.norm_stats["actions"]).astype(np.float32) diff --git a/dimos/imitation/policy/abc/python/abc_minimal/preprocess.py b/dimos/imitation/policy/abc/python/abc_minimal/preprocess.py new file mode 100644 index 0000000000..a1306afc37 --- /dev/null +++ b/dimos/imitation/policy/abc/python/abc_minimal/preprocess.py @@ -0,0 +1,71 @@ +"""Shared state/action normalization and image preprocessing.""" + +import json +from pathlib import Path + +import numpy as np +import torch +import torch.nn.functional as F + + +IMAGENET_MEAN = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1) +IMAGENET_STD = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1) + + +def parse_norm_stats(raw): + stats = raw.get("norm_stats", raw) + if "state" not in stats and "actions" not in stats: + key = "xdof" if "xdof" in stats else next(iter(stats)) + stats = stats[key] + return { + key: {k: np.asarray(v, dtype=np.float32) for k, v in stats[key].items()} + for key in ("state", "actions") + } + + +def load_norm_stats(path): + return parse_norm_stats(json.loads(Path(path).read_text())) + + +def normalize(x, stats): + return (x - stats["mean"]) / (stats["std"] + 1e-6) + + +def unnormalize(x, stats): + return x * (stats["std"] + 1e-6) + stats["mean"] + + +def resize_with_pad(img_hwc, target_h=224, target_w=224): + h, w, _ = img_hwc.shape + if (h, w) == (target_h, target_w): + return img_hwc + ratio = max(w / target_w, h / target_h) + new_h = max(1, int(round(h / ratio))) + new_w = max(1, int(round(w / ratio))) + resized = F.interpolate( + img_hwc.permute(2, 0, 1).unsqueeze(0), + size=(new_h, new_w), + mode="bilinear", + align_corners=False, + antialias=True, + ).squeeze(0) + pad_h0 = (target_h - new_h) // 2 + pad_h1 = target_h - new_h - pad_h0 + pad_w0 = (target_w - new_w) // 2 + pad_w1 = target_w - new_w - pad_w0 + padded = F.pad(resized, (pad_w0, pad_w1, pad_h0, pad_h1), value=0) + return padded.permute(1, 2, 0) + + +def imagenet_normalize(img_chw): + mean = IMAGENET_MEAN.to(device=img_chw.device, dtype=img_chw.dtype) + std = IMAGENET_STD.to(device=img_chw.device, dtype=img_chw.dtype) + return (img_chw - mean) / (std + 1e-6) + + +def resize_pad_normalize(img_chw, target_h=224, target_w=224): + x = torch.as_tensor(img_chw).float() + if x.max() > 1.0: + x = x / 255.0 + x = resize_with_pad(x.permute(1, 2, 0), target_h, target_w).permute(2, 0, 1) + return imagenet_normalize(x) diff --git a/dimos/imitation/policy/abc/python/dimos_abc/__init__.py b/dimos/imitation/policy/abc/python/dimos_abc/__init__.py new file mode 100644 index 0000000000..e908044542 --- /dev/null +++ b/dimos/imitation/policy/abc/python/dimos_abc/__init__.py @@ -0,0 +1 @@ +"""DimOS adapter for the vendored ABC-DiT inference code.""" diff --git a/dimos/imitation/policy/abc/python/dimos_abc/runtime.py b/dimos/imitation/policy/abc/python/dimos_abc/runtime.py new file mode 100644 index 0000000000..1c8d21e2d6 --- /dev/null +++ b/dimos/imitation/policy/abc/python/dimos_abc/runtime.py @@ -0,0 +1,208 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Amazon ABC-DiT adapter for the shared DimOS rollout runtime.""" + +from __future__ import annotations + +from collections.abc import Mapping +import json +from pathlib import Path +from types import SimpleNamespace +from typing import Any, cast + +from abc_minimal.config import ClipConfig, DiTConfig +from abc_minimal.dit import CLIPTextEmbedder, DiTPolicy, load_pretrained +from abc_minimal.fast_inference import FastInferenceGraph +from abc_minimal.preprocess import normalize, parse_norm_stats, resize_pad_normalize, unnormalize +import numpy as np +from numpy.typing import NDArray +import torch + +from dimos.imitation.policy.abc.module import AbcPolicyConfig, DualOpenYamAbcPolicy +from dimos.imitation.policy.backend import PolicyBackendInfo +from dimos.imitation.policy.runtime import declare_policy_runtime +from dimos.imitation.profile import ImageSource, JointPositionSource, PolicyIOProfile + +torch.set_float32_matmul_precision("high") + + +class AbcBackend: + """Run the released 14-D, three-camera ABC-DiT policy in process.""" + + def __init__(self, config: AbcPolicyConfig) -> None: + self._rollout_config = config + self.config = SimpleNamespace(model=DiTConfig()) + self.device = torch.device("cpu") + self.diffusion_steps = config.diffusion_steps + self.model: DiTPolicy | None = None + self.embedder: CLIPTextEmbedder | None = None + self.task_vec = torch.empty(0) + self.norm_stats: dict[str, Any] = {} + self._task = "" + self._fast_graph: FastInferenceGraph | None = None + + def load(self, profile: PolicyIOProfile) -> PolicyBackendInfo: + _validate_profile(profile, self.config.model) + checkpoint = Path(self._rollout_config.artifact).expanduser().resolve() + if not checkpoint.is_file(): + raise FileNotFoundError(f"ABC checkpoint does not exist: {checkpoint}") + device_name = self._rollout_config.device or ( + "cuda" if torch.cuda.is_available() else "cpu" + ) + self.device = torch.device(device_name) + if self.device.type == "cuda" and not torch.cuda.is_available(): + raise RuntimeError(f"ABC requested device {device_name!r}, but CUDA is unavailable") + + self.model = DiTPolicy(self.config.model).to(self.device) + checkpoint_data = load_pretrained(self.model, checkpoint) + self.model.eval() + self.norm_stats = _resolve_norm_stats( + checkpoint_data, + self._rollout_config.norm_stats_path, + ) + _validate_norm_stats(self.norm_stats, self.config.model) + self.embedder = CLIPTextEmbedder(ClipConfig(), device=self.device) + self._set_task(self._rollout_config.task) + return PolicyBackendInfo( + name="abc", + chunk_length=self.config.model.chunk_length, + preferred_execution_steps=15, + ) + + def reset(self) -> None: + """ABC-DiT has no recurrent or queued model state.""" + + @torch.no_grad() + def predict( + self, + observations: Mapping[str, NDArray[Any]], + task: str, + ) -> NDArray[np.float32]: + model = self._require_model() + if task != self._task: + if self._fast_graph is not None: + raise ValueError("task text cannot change after ABC fast inference is captured") + self._set_task(task) + images: dict[str, NDArray[Any]] = { + camera: np.asarray(observations[camera]).transpose(2, 0, 1) + for camera in self.config.model.camera_keys + } + obs: dict[str, Any] = { + "state": np.asarray(observations["state"], dtype=np.float32), + "images": images, + } + if self._rollout_config.fast_inference and self.device.type == "cuda": + if self._fast_graph is None: + self._enable_fast_inference(obs) + assert self._fast_graph is not None + return self._fast_graph.infer(obs, noise=None) + + state = normalize(obs["state"], self.norm_stats["state"]) + batch = { + "state": torch.from_numpy(state[None]).float().to(self.device), + "actions": torch.zeros( + 1, + self.config.model.chunk_length, + self.config.model.action_dim, + device=self.device, + ), + "images": { + camera: resize_pad_normalize(obs["images"][camera]).unsqueeze(0).to(self.device) + for camera in self.config.model.camera_keys + }, + "task_vec_clip": self.task_vec, + } + actions = model.sample_actions(batch, num_steps=self.diffusion_steps) + result = actions[0].float().detach().cpu().numpy() + return np.asarray( + unnormalize(result, self.norm_stats["actions"]), + dtype=np.float32, + ) + + def _set_task(self, task: str) -> None: + if self.embedder is None: + raise RuntimeError("ABC text embedder is not loaded") + self.task_vec = self.embedder.encode([task]).to(self.device) + self._task = task + + def _enable_fast_inference(self, observation: dict[str, Any]) -> None: + model = self._require_model() + torch.backends.cuda.matmul.allow_tf32 = True + torch.backends.cudnn.allow_tf32 = True + model.to(torch.bfloat16) + model.img_backbone.set_bfloat16(True) + self.task_vec = self.task_vec.to(device=self.device, dtype=torch.bfloat16) + model.predict_velocity = torch.compile( # type: ignore[method-assign] + model.predict_velocity, + dynamic=False, + mode="max-autotune-no-cudagraphs", + ) + graph = FastInferenceGraph(self) + graph.capture(observation, warmup_noise=None, replay_warmups=24) + self._fast_graph = graph + + def _require_model(self) -> DiTPolicy: + if self.model is None: + raise RuntimeError("ABC backend is not loaded") + return self.model + + +AbcPolicyRuntime = declare_policy_runtime( + "AbcPolicyRuntime", + __name__, + DualOpenYamAbcPolicy, + AbcBackend, +) + + +def _validate_profile(profile: PolicyIOProfile, model: DiTConfig) -> None: + expected = {*model.camera_keys, "state"} + if set(profile.observations) != expected: + raise ValueError(f"released ABC checkpoint requires observation keys {sorted(expected)}") + for key in model.camera_keys: + if not isinstance(profile.observations[key], ImageSource): + raise TypeError(f"ABC observation {key!r} must be an image") + state = profile.observations["state"] + if not isinstance(state, JointPositionSource) or len(state.joints) != model.state_dim: + raise ValueError(f"ABC state must contain {model.state_dim} joints") + if profile.action.key != "actions": + raise ValueError("released ABC checkpoint requires the 'actions' output key") + if len(profile.action.demonstration.joints) != model.action_dim: + raise ValueError(f"ABC actions must contain {model.action_dim} joints") + + +def _resolve_norm_stats( + checkpoint: dict[str, Any], + override: str | None, +) -> dict[str, Any]: + if override is not None: + raw = json.loads(Path(override).expanduser().read_text()) + elif checkpoint.get("norm_stats") is not None: + raw = checkpoint["norm_stats"] + else: + raise ValueError("ABC checkpoint has no norm_stats; set norm_stats_path") + return cast("dict[str, Any]", parse_norm_stats(raw)) + + +def _validate_norm_stats(stats: dict[str, Any], model: DiTConfig) -> None: + for key, width in (("state", model.state_dim), ("actions", model.action_dim)): + for statistic in ("mean", "std"): + value = np.asarray(stats[key][statistic]) + if value.shape != (width,): + raise ValueError( + f"ABC {key} {statistic} shape {value.shape} does not match {(width,)}" + ) + if not np.all(np.isfinite(value)): + raise ValueError(f"ABC {key} {statistic} contains non-finite values") diff --git a/dimos/imitation/policy/abc/python/dimos_abc/runtime_tests.py b/dimos/imitation/policy/abc/python/dimos_abc/runtime_tests.py new file mode 100644 index 0000000000..0d4ced30be --- /dev/null +++ b/dimos/imitation/policy/abc/python/dimos_abc/runtime_tests.py @@ -0,0 +1,40 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from abc_minimal.config import DiTConfig +import numpy as np +import pytest + +from dimos.experimental.isolated_python.bootstrap import validate_runtime +from dimos.imitation.policy.abc.module import DualOpenYamAbcPolicy +from dimos.robot.manipulators.dual_openyam.learning import DUAL_OPENYAM_ABC_IO +from dimos_abc.runtime import AbcPolicyRuntime, _validate_norm_stats, _validate_profile + + +def test_generated_runtime_implements_the_host_contract() -> None: + validate_runtime(DualOpenYamAbcPolicy, AbcPolicyRuntime) + + +def test_released_abc_profile_matches_vendored_model_contract() -> None: + _validate_profile(DUAL_OPENYAM_ABC_IO, DiTConfig()) + + +def test_norm_stats_require_exact_released_dimensions() -> None: + stats = { + "state": {"mean": np.zeros(14), "std": np.ones(14)}, + "actions": {"mean": np.zeros(13), "std": np.ones(13)}, + } + + with pytest.raises(ValueError, match="actions mean shape"): + _validate_norm_stats(stats, DiTConfig()) diff --git a/dimos/imitation/policy/abc/python/pyproject.toml b/dimos/imitation/policy/abc/python/pyproject.toml new file mode 100644 index 0000000000..62f4f57e07 --- /dev/null +++ b/dimos/imitation/policy/abc/python/pyproject.toml @@ -0,0 +1,56 @@ +[build-system] +requires = ["setuptools>=70"] +build-backend = "setuptools.build_meta" + +[project] +name = "dimos-abc-runtime" +version = "0.1.0" +requires-python = ">=3.12,<3.13" +dependencies = [ + "ftfy>=6.3,<7", + "numpy>=2,<3", + "regex>=2025.7", + "torch==2.11.0+cu128", +] + +[[tool.uv.index]] +name = "pytorch-cu128" +url = "https://download.pytorch.org/whl/cu128" +explicit = true + +[tool.uv.sources] +torch = { index = "pytorch-cu128" } + +[dependency-groups] +tests = [ + "mypy==1.19.0", + "pytest==8.3.5", + "pytest-mock>=3.14", +] + +[tool.uv] +default-groups = [] + +[tool.setuptools.packages.find] +where = ["."] +include = ["abc_minimal*", "dimos_abc*"] + +[tool.pytest.ini_options] +testpaths = ["dimos_abc"] +python_files = ["*_tests.py"] + +[tool.mypy] +files = ["dimos_abc/"] +python_version = "3.12" +strict = true +explicit_package_bases = true +mypy_path = "../../../../../" +untyped_calls_exclude = ["abc_minimal"] + +[[tool.mypy.overrides]] +module = ["dimos", "dimos.*"] +follow_imports = "skip" + +[[tool.mypy.overrides]] +module = ["abc_minimal", "abc_minimal.*"] +follow_untyped_imports = true diff --git a/dimos/imitation/policy/abc/python/uv.lock b/dimos/imitation/policy/abc/python/uv.lock new file mode 100644 index 0000000000..3ddbaf292e --- /dev/null +++ b/dimos/imitation/policy/abc/python/uv.lock @@ -0,0 +1,569 @@ +version = 1 +revision = 3 +requires-python = "==3.12.*" + +[[package]] +name = "colorama" +version = "0.4.6" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/d8/53/6f443c9a4a8358a93a6792e2acffb9d9d5cb0a5cfd8802644b7b1c9a02e4/colorama-0.4.6.tar.gz", hash = "sha256:08695f5cb7ed6e0531a20572697297273c47b8cae5a63ffc6d6ed5c201be6e44", size = 27697, upload-time = "2022-10-25T02:36:22.414Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d1/d6/3965ed04c63042e047cb6a3e6ed1a63a35087b6a609aa3a15ed8ac56c221/colorama-0.4.6-py2.py3-none-any.whl", hash = "sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6", size = 25335, upload-time = "2022-10-25T02:36:20.889Z" }, +] + +[[package]] +name = "cuda-bindings" +version = "12.9.7" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "cuda-pathfinder" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/32/45/557d4ed1fa54f0c7db8aee083229f624990d69f7d00f55477eed5c7e169a/cuda_bindings-12.9.7-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:0666d3c082ef8f4b2d670950589373550e9f3bf564d635dd883f24a0b40402ff", size = 7071026, upload-time = "2026-05-27T18:44:13.356Z" }, + { url = "https://files.pythonhosted.org/packages/91/97/e3c6e58ece26a053419ba0a18444b5443cfc64451bbf37f84e8143b8bdca/cuda_bindings-12.9.7-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:4c7ef48c5e13ae90f3b2ecfb72f8e99ac43c8f4c43e67e1325b8aae331453687", size = 7611059, upload-time = "2026-05-27T18:44:15.252Z" }, +] + +[[package]] +name = "cuda-pathfinder" +version = "1.8.1" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/9d/e6/22df83f82f9bc26cb1c42265cf14d34d4908dba2a0f261bd7b28244acb00/cuda_pathfinder-1.8.1-py3-none-any.whl", hash = "sha256:ae0137ff9e56ea97499bcbf54f5f2778ec25f3266715ac86da192a795af982a8", size = 62552, upload-time = "2026-09-02T16:55:28.64Z" }, +] + +[[package]] +name = "cuda-toolkit" +version = "12.8.1" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d4/c8/7dce3a0b15b42a3b58e7d96eb22a687d3bf2c44e01d149a6874629cd9938/cuda_toolkit-12.8.1-py2.py3-none-any.whl", hash = "sha256:adc7906af4ecbf9a352f9dca5734eceb21daec281ccfcf5675e1d2f724fc2cba", size = 2283, upload-time = "2025-08-13T02:03:07.842Z" }, +] + +[package.optional-dependencies] +cublas = [ + { name = "nvidia-cublas-cu12", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, +] +cudart = [ + { name = "nvidia-cuda-runtime-cu12", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, +] +cufft = [ + { name = "nvidia-cufft-cu12", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, +] +cufile = [ + { name = "nvidia-cufile-cu12", marker = "sys_platform == 'linux'" }, +] +cupti = [ + { name = "nvidia-cuda-cupti-cu12", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, +] +curand = [ + { name = "nvidia-curand-cu12", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, +] +cusolver = [ + { name = "nvidia-cusolver-cu12", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, +] +cusparse = [ + { name = "nvidia-cusparse-cu12", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, +] +nvjitlink = [ + { name = "nvidia-nvjitlink-cu12", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, +] +nvrtc = [ + { name = "nvidia-cuda-nvrtc-cu12", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, +] +nvtx = [ + { name = "nvidia-nvtx-cu12", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, +] + +[[package]] +name = "dimos-abc-runtime" +version = "0.1.0" +source = { editable = "." } +dependencies = [ + { name = "ftfy" }, + { name = "numpy" }, + { name = "regex" }, + { name = "torch" }, +] + +[package.dev-dependencies] +tests = [ + { name = "mypy" }, + { name = "pytest" }, + { name = "pytest-mock" }, +] + +[package.metadata] +requires-dist = [ + { name = "ftfy", specifier = ">=6.3,<7" }, + { name = "numpy", specifier = ">=2,<3" }, + { name = "regex", specifier = ">=2025.7" }, + { name = "torch", specifier = "==2.11.0+cu128", index = "https://download.pytorch.org/whl/cu128" }, +] + +[package.metadata.requires-dev] +tests = [ + { name = "mypy", specifier = "==1.19.0" }, + { name = "pytest", specifier = "==8.3.5" }, + { name = "pytest-mock", specifier = ">=3.14" }, +] + +[[package]] +name = "filelock" +version = "3.32.5" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/0a/a0/50c2c0ce5e74d7721bbb1b19a26ebd339aac5878553a6e35308c2f31f935/filelock-3.32.5.tar.gz", hash = "sha256:f6a6a28f743f9b95ce19db5abe0f376f75eb56517dff21e1a4751e2657d3e83d", size = 222838, upload-time = "2026-08-31T18:56:34.729Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/36/d2/b70a31e13d04456d28493f31d2aa087e99eeb2767ef0293b2625727ccb8c/filelock-3.32.5-py3-none-any.whl", hash = "sha256:142cd9fa77a872c5e78c62329a0d15278fadc686eb89e760017968961a4fd6b2", size = 100003, upload-time = "2026-08-31T18:56:33.078Z" }, +] + +[[package]] +name = "fsspec" +version = "2026.7.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/00/78/f34251dadb8f3921264a1d9b8946f5e542014ee2614b285261b4e40e6775/fsspec-2026.7.0.tar.gz", hash = "sha256:c803c40f4cf860b49dea58ee3e1c33cb9c790520e233537e1340049f89b82a88", size = 317040, upload-time = "2026-07-28T16:34:51.052Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fd/3c/6a2bf344106328fd04963664a60b9bb6496fc25df8e962fcdc1367285fb9/fsspec-2026.7.0-py3-none-any.whl", hash = "sha256:b57ddbafedfaef7018c1ecab32aa200a9d7ca26b77965f64e48b70061249d279", size = 206583, upload-time = "2026-07-28T16:34:49.538Z" }, +] + +[[package]] +name = "ftfy" +version = "6.3.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "wcwidth" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/a5/d3/8650919bc3c7c6e90ee3fa7fd618bf373cbbe55dff043bd67353dbb20cd8/ftfy-6.3.1.tar.gz", hash = "sha256:9b3c3d90f84fb267fe64d375a07b7f8912d817cf86009ae134aa03e1819506ec", size = 308927, upload-time = "2024-10-26T00:50:35.149Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ab/6e/81d47999aebc1b155f81eca4477a616a70f238a2549848c38983f3c22a82/ftfy-6.3.1-py3-none-any.whl", hash = "sha256:7c70eb532015cd2f9adb53f101fb6c7945988d023a085d127d1573dc49dd0083", size = 44821, upload-time = "2024-10-26T00:50:33.425Z" }, +] + +[[package]] +name = "iniconfig" +version = "2.3.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/72/34/14ca021ce8e5dfedc35312d08ba8bf51fdd999c576889fc2c24cb97f4f10/iniconfig-2.3.0.tar.gz", hash = "sha256:c76315c77db068650d49c5b56314774a7804df16fee4402c1f19d6d15d8c4730", size = 20503, upload-time = "2025-10-18T21:55:43.219Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/cb/b1/3846dd7f199d53cb17f49cba7e651e9ce294d8497c8c150530ed11865bb8/iniconfig-2.3.0-py3-none-any.whl", hash = "sha256:f631c04d2c48c52b84d0d0549c99ff3859c98df65b3101406327ecc7d53fbf12", size = 7484, upload-time = "2025-10-18T21:55:41.639Z" }, +] + +[[package]] +name = "jinja2" +version = "3.1.6" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "markupsafe" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/df/bf/f7da0350254c0ed7c72f3e33cef02e048281fec7ecec5f032d4aac52226b/jinja2-3.1.6.tar.gz", hash = "sha256:0137fb05990d35f1275a587e9aee6d56da821fc83491a0fb838183be43f66d6d", size = 245115, upload-time = "2025-03-05T20:05:02.478Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/62/a1/3d680cbfd5f4b8f15abc1d571870c5fc3e594bb582bc3b64ea099db13e56/jinja2-3.1.6-py3-none-any.whl", hash = "sha256:85ece4451f492d0c13c5dd7c13a64681a86afae63a5f347908daf103ce6d2f67", size = 134899, upload-time = "2025-03-05T20:05:00.369Z" }, +] + +[[package]] +name = "librt" +version = "0.15.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/36/9b/356320fbae2ac8467e21c5e73e1389c80468e4998c62cc7d3536cc51b614/librt-0.15.0.tar.gz", hash = "sha256:4e66cbe84437497d951b799d3e1551291b6fb3d643820a7014b3655d57a59162", size = 214338, upload-time = "2026-08-07T10:49:42.663Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ba/39/99c25030e782bdfb7a21be8c05254806a2e4bbb05c8d50c2a2130acbfa05/librt-0.15.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:e87bc679f86a99aa3b26e3c78eeb821a247c9a28eae48eaafcc32c3bf4c3bb9e", size = 151021, upload-time = "2026-08-07T10:47:00.057Z" }, + { url = "https://files.pythonhosted.org/packages/14/43/f4b1bd1b2888798a1409808889a25ea1ba49eaabce7d681ed27734c2df9d/librt-0.15.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:71599e011ac880e8e45d46047d714871894c7d4ab6f25626f8d4f89da21f368d", size = 155267, upload-time = "2026-08-07T10:47:01.311Z" }, + { url = "https://files.pythonhosted.org/packages/0c/db/3ad9c965c72f1e1d6beeec44ec10a54e17be8ae042fbb4baade16cbadced/librt-0.15.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c802434092b769b1d613ed2e13fac15fbfce1934a74bd10283b03c0fae231cd1", size = 503136, upload-time = "2026-08-07T10:47:02.45Z" }, + { url = "https://files.pythonhosted.org/packages/4b/07/5888a6d76acd62ebce66c61b74d94e9370b9c32929f111e487bb6546f8ed/librt-0.15.0-cp312-cp312-manylinux2014_i686.manylinux_2_17_i686.manylinux_2_28_i686.whl", hash = "sha256:5500eeae393a184d14e1f35645962c27129d20c81afa4069e6ef826ebc2b3aaa", size = 496670, upload-time = "2026-08-07T10:47:03.675Z" }, + { url = "https://files.pythonhosted.org/packages/29/39/ab57cc2f5b276156da02bb7f5a8921bada1cb1993ffec99acf811c602c23/librt-0.15.0-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:6ecfc32dfb46fb7b565bcd6abf9412acf978775a998273d22888a6d7953730dd", size = 513688, upload-time = "2026-08-07T10:47:04.981Z" }, + { url = "https://files.pythonhosted.org/packages/a7/b9/bdbb0b648b5c2befb031f4c6f3b1dd857415e8fb492a25a3c764a6681e6c/librt-0.15.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:89cc46cfd15022e35084355478c9ac809d90b1152222706ac9a7655ec21df6fa", size = 531904, upload-time = "2026-08-07T10:47:06.211Z" }, + { url = "https://files.pythonhosted.org/packages/93/26/473c2e4b6c104e9e58e27ce95fc8005c8bd4fc36cae4f254371125a92db8/librt-0.15.0-cp312-cp312-manylinux_2_34_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:d5f51401d102c885b9ca509e62c79b1dbff286e1b9b047fde6f763780789356d", size = 524427, upload-time = "2026-08-07T10:47:07.592Z" }, + { url = "https://files.pythonhosted.org/packages/26/60/03b3abb82b41714671b907bf6989b228e31e6a8af52dec82b5b0728dc250/librt-0.15.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:cc30523e3f1a23fb7511cc659834a0d01a1042bb9de359bc1c131cc4ec6c9656", size = 543155, upload-time = "2026-08-07T10:47:08.866Z" }, + { url = "https://files.pythonhosted.org/packages/f2/0e/9bb1f0a4affbd0a1888f4f79dc03ed2a299d9a2c26c59ab2a97dcbf11903/librt-0.15.0-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:59fe030d8ae4a57e3fb7756bf35a858de74e04066fc8555c53d0af979132af81", size = 546890, upload-time = "2026-08-07T10:47:10.327Z" }, + { url = "https://files.pythonhosted.org/packages/dc/84/6937a280d461f7de6e031ffb02edc2b7c3c90d49d630565ce8ff27cbc5f2/librt-0.15.0-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:5a6526a2a956bbb1e4ae3568c82e650fc99119c66bb011ea60715744955a2b4d", size = 555163, upload-time = "2026-08-07T10:47:11.798Z" }, + { url = "https://files.pythonhosted.org/packages/bc/95/2a2853c1ee014bf102116e7f897a04beeaeb2461b45b79af98bdfb95f1ef/librt-0.15.0-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:85ea21ec6730194d67156b0e0b5430ccb1d61f8b8b907e39b37f9812b74a13f0", size = 535812, upload-time = "2026-08-07T10:47:13.279Z" }, + { url = "https://files.pythonhosted.org/packages/c9/4c/cf9601c1b4c5f09280acd5d83abdb2e68527a2be8257136eb42304218622/librt-0.15.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:1e47b8ba865d7ede071a91a7163073bbaeb72541f1ef8a07d512c45c7b5007f2", size = 573688, upload-time = "2026-08-07T10:47:14.727Z" }, + { url = "https://files.pythonhosted.org/packages/47/6d/9ac7cbec46189a7625af4b5acbd25f10d827f4141b2002181848c8418923/librt-0.15.0-cp312-cp312-win32.whl", hash = "sha256:a5207ec414d1c4a2a7231b2086970dc036f94293cdf338190984958a013a42f1", size = 106138, upload-time = "2026-08-07T10:47:15.973Z" }, + { url = "https://files.pythonhosted.org/packages/38/d0/2ae99c83be86ce23f925ac1aeeedc777e97f427c4a8d190c70d0a16e9a87/librt-0.15.0-cp312-cp312-win_amd64.whl", hash = "sha256:73b30cfa976659b3917c8f6153bdb0591c6a9ec6583599fd24a689b690622022", size = 126974, upload-time = "2026-08-07T10:47:17.049Z" }, + { url = "https://files.pythonhosted.org/packages/5d/ef/dd24f9635c730b86b87587967dda7516b1845e8b17684603d31607fed598/librt-0.15.0-cp312-cp312-win_arm64.whl", hash = "sha256:a54cf9e0ef47b96af580849db5471142200568ce1e02cbf416addab551369570", size = 112292, upload-time = "2026-08-07T10:47:18.222Z" }, +] + +[[package]] +name = "markupsafe" +version = "3.0.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/7e/99/7690b6d4034fffd95959cbe0c02de8deb3098cc577c67bb6a24fe5d7caa7/markupsafe-3.0.3.tar.gz", hash = "sha256:722695808f4b6457b320fdc131280796bdceb04ab50fe1795cd540799ebe1698", size = 80313, upload-time = "2025-09-27T18:37:40.426Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5a/72/147da192e38635ada20e0a2e1a51cf8823d2119ce8883f7053879c2199b5/markupsafe-3.0.3-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:d53197da72cc091b024dd97249dfc7794d6a56530370992a5e1a08983ad9230e", size = 11615, upload-time = "2025-09-27T18:36:30.854Z" }, + { url = "https://files.pythonhosted.org/packages/9a/81/7e4e08678a1f98521201c3079f77db69fb552acd56067661f8c2f534a718/markupsafe-3.0.3-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:1872df69a4de6aead3491198eaf13810b565bdbeec3ae2dc8780f14458ec73ce", size = 12020, upload-time = "2025-09-27T18:36:31.971Z" }, + { url = "https://files.pythonhosted.org/packages/1e/2c/799f4742efc39633a1b54a92eec4082e4f815314869865d876824c257c1e/markupsafe-3.0.3-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:3a7e8ae81ae39e62a41ec302f972ba6ae23a5c5396c8e60113e9066ef893da0d", size = 24332, upload-time = "2025-09-27T18:36:32.813Z" }, + { url = "https://files.pythonhosted.org/packages/3c/2e/8d0c2ab90a8c1d9a24f0399058ab8519a3279d1bd4289511d74e909f060e/markupsafe-3.0.3-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d6dd0be5b5b189d31db7cda48b91d7e0a9795f31430b7f271219ab30f1d3ac9d", size = 22947, upload-time = "2025-09-27T18:36:33.86Z" }, + { url = "https://files.pythonhosted.org/packages/2c/54/887f3092a85238093a0b2154bd629c89444f395618842e8b0c41783898ea/markupsafe-3.0.3-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:94c6f0bb423f739146aec64595853541634bde58b2135f27f61c1ffd1cd4d16a", size = 21962, upload-time = "2025-09-27T18:36:35.099Z" }, + { url = "https://files.pythonhosted.org/packages/c9/2f/336b8c7b6f4a4d95e91119dc8521402461b74a485558d8f238a68312f11c/markupsafe-3.0.3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:be8813b57049a7dc738189df53d69395eba14fb99345e0a5994914a3864c8a4b", size = 23760, upload-time = "2025-09-27T18:36:36.001Z" }, + { url = "https://files.pythonhosted.org/packages/32/43/67935f2b7e4982ffb50a4d169b724d74b62a3964bc1a9a527f5ac4f1ee2b/markupsafe-3.0.3-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:83891d0e9fb81a825d9a6d61e3f07550ca70a076484292a70fde82c4b807286f", size = 21529, upload-time = "2025-09-27T18:36:36.906Z" }, + { url = "https://files.pythonhosted.org/packages/89/e0/4486f11e51bbba8b0c041098859e869e304d1c261e59244baa3d295d47b7/markupsafe-3.0.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:77f0643abe7495da77fb436f50f8dab76dbc6e5fd25d39589a0f1fe6548bfa2b", size = 23015, upload-time = "2025-09-27T18:36:37.868Z" }, + { url = "https://files.pythonhosted.org/packages/2f/e1/78ee7a023dac597a5825441ebd17170785a9dab23de95d2c7508ade94e0e/markupsafe-3.0.3-cp312-cp312-win32.whl", hash = "sha256:d88b440e37a16e651bda4c7c2b930eb586fd15ca7406cb39e211fcff3bf3017d", size = 14540, upload-time = "2025-09-27T18:36:38.761Z" }, + { url = "https://files.pythonhosted.org/packages/aa/5b/bec5aa9bbbb2c946ca2733ef9c4ca91c91b6a24580193e891b5f7dbe8e1e/markupsafe-3.0.3-cp312-cp312-win_amd64.whl", hash = "sha256:26a5784ded40c9e318cfc2bdb30fe164bdb8665ded9cd64d500a34fb42067b1c", size = 15105, upload-time = "2025-09-27T18:36:39.701Z" }, + { url = "https://files.pythonhosted.org/packages/e5/f1/216fc1bbfd74011693a4fd837e7026152e89c4bcf3e77b6692fba9923123/markupsafe-3.0.3-cp312-cp312-win_arm64.whl", hash = "sha256:35add3b638a5d900e807944a078b51922212fb3dedb01633a8defc4b01a3c85f", size = 13906, upload-time = "2025-09-27T18:36:40.689Z" }, +] + +[[package]] +name = "mpmath" +version = "1.3.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/e0/47/dd32fa426cc72114383ac549964eecb20ecfd886d1e5ccf5340b55b02f57/mpmath-1.3.0.tar.gz", hash = "sha256:7a28eb2a9774d00c7bc92411c19a89209d5da7c4c9a9e227be8330a23a25b91f", size = 508106, upload-time = "2023-03-07T16:47:11.061Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/43/e3/7d92a15f894aa0c9c4b49b8ee9ac9850d6e63b03c9c32c0367a13ae62209/mpmath-1.3.0-py3-none-any.whl", hash = "sha256:a0b2b9fe80bbcd81a6647ff13108738cfb482d481d826cc0e02f5b35e5c88d2c", size = 536198, upload-time = "2023-03-07T16:47:09.197Z" }, +] + +[[package]] +name = "mypy" +version = "1.19.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "librt" }, + { name = "mypy-extensions" }, + { name = "pathspec" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/f9/b5/b58cdc25fadd424552804bf410855d52324183112aa004f0732c5f6324cf/mypy-1.19.0.tar.gz", hash = "sha256:f6b874ca77f733222641e5c46e4711648c4037ea13646fd0cdc814c2eaec2528", size = 3579025, upload-time = "2025-11-28T15:49:01.26Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/11/7e/1afa8fb188b876abeaa14460dc4983f909aaacaa4bf5718c00b2c7e0b3d5/mypy-1.19.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:0fb3115cb8fa7c5f887c8a8d81ccdcb94cff334684980d847e5a62e926910e1d", size = 13207728, upload-time = "2025-11-28T15:46:26.463Z" }, + { url = "https://files.pythonhosted.org/packages/b2/13/f103d04962bcbefb1644f5ccb235998b32c337d6c13145ea390b9da47f3e/mypy-1.19.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:f3e19e3b897562276bb331074d64c076dbdd3e79213f36eed4e592272dabd760", size = 12202945, upload-time = "2025-11-28T15:48:49.143Z" }, + { url = "https://files.pythonhosted.org/packages/e4/93/a86a5608f74a22284a8ccea8592f6e270b61f95b8588951110ad797c2ddd/mypy-1.19.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:b9d491295825182fba01b6ffe2c6fe4e5a49dbf4e2bb4d1217b6ced3b4797bc6", size = 12718673, upload-time = "2025-11-28T15:47:37.193Z" }, + { url = "https://files.pythonhosted.org/packages/3d/58/cf08fff9ced0423b858f2a7495001fda28dc058136818ee9dffc31534ea9/mypy-1.19.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:6016c52ab209919b46169651b362068f632efcd5eb8ef9d1735f6f86da7853b2", size = 13608336, upload-time = "2025-11-28T15:48:32.625Z" }, + { url = "https://files.pythonhosted.org/packages/64/ed/9c509105c5a6d4b73bb08733102a3ea62c25bc02c51bca85e3134bf912d3/mypy-1.19.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:f188dcf16483b3e59f9278c4ed939ec0254aa8a60e8fc100648d9ab5ee95a431", size = 13833174, upload-time = "2025-11-28T15:45:48.091Z" }, + { url = "https://files.pythonhosted.org/packages/cd/71/01939b66e35c6f8cb3e6fdf0b657f0fd24de2f8ba5e523625c8e72328208/mypy-1.19.0-cp312-cp312-win_amd64.whl", hash = "sha256:0e3c3d1e1d62e678c339e7ade72746a9e0325de42cd2cccc51616c7b2ed1a018", size = 10112208, upload-time = "2025-11-28T15:46:41.702Z" }, + { url = "https://files.pythonhosted.org/packages/09/0e/fe228ed5aeab470c6f4eb82481837fadb642a5aa95cc8215fd2214822c10/mypy-1.19.0-py3-none-any.whl", hash = "sha256:0c01c99d626380752e527d5ce8e69ffbba2046eb8a060db0329690849cf9b6f9", size = 2469714, upload-time = "2025-11-28T15:45:33.22Z" }, +] + +[[package]] +name = "mypy-extensions" +version = "1.1.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/a2/6e/371856a3fb9d31ca8dac321cda606860fa4548858c0cc45d9d1d4ca2628b/mypy_extensions-1.1.0.tar.gz", hash = "sha256:52e68efc3284861e772bbcd66823fde5ae21fd2fdb51c62a211403730b916558", size = 6343, upload-time = "2025-04-22T14:54:24.164Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/79/7b/2c79738432f5c924bef5071f933bcc9efd0473bac3b4aa584a6f7c1c8df8/mypy_extensions-1.1.0-py3-none-any.whl", hash = "sha256:1be4cccdb0f2482337c4743e60421de3a356cd97508abadd57d47403e94f5505", size = 4963, upload-time = "2025-04-22T14:54:22.983Z" }, +] + +[[package]] +name = "networkx" +version = "3.6.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/6a/51/63fe664f3908c97be9d2e4f1158eb633317598cfa6e1fc14af5383f17512/networkx-3.6.1.tar.gz", hash = "sha256:26b7c357accc0c8cde558ad486283728b65b6a95d85ee1cd66bafab4c8168509", size = 2517025, upload-time = "2025-12-08T17:02:39.908Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/9e/c9/b2622292ea83fbb4ec318f5b9ab867d0a28ab43c5717bb85b0a5f6b3b0a4/networkx-3.6.1-py3-none-any.whl", hash = "sha256:d47fbf302e7d9cbbb9e2555a0d267983d2aa476bac30e90dfbe5669bd57f3762", size = 2068504, upload-time = "2025-12-08T17:02:38.159Z" }, +] + +[[package]] +name = "numpy" +version = "2.5.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/9a/80/db0b4559e57ec36362bedbb05530a87fafbcb6067708c946967a41d449e7/numpy-2.5.2.tar.gz", hash = "sha256:d482d171c406ae88c5b19cad3b6a1c4c5209f886ab74bc44c2c865c23f52d860", size = 20773161, upload-time = "2026-08-09T13:48:27.962Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/69/72/dccb0aaf40972777283303919f613964227266d0c13adebb79ac124f1c3e/numpy-2.5.2-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:14e373cfc6387177e8409dac3c7159be8eb05cd77096cd7c950268b86f62831c", size = 16891693, upload-time = "2026-08-09T13:44:51.702Z" }, + { url = "https://files.pythonhosted.org/packages/60/2e/b5aee50a1f74ac815cf8331812cb8251e29024025de462e0c047641c614c/numpy-2.5.2-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:4bbd96c833ecc8cc069ce518078fc8c60cb9cbfb0fea5b7a803ad65035596d03", size = 11903109, upload-time = "2026-08-09T13:44:55.501Z" }, + { url = "https://files.pythonhosted.org/packages/f3/f4/29e78102a80601cf034d4e9767022cffeca2c3b4c926e1754572ca95593d/numpy-2.5.2-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:6e8172ddfcf5cf74b811d372b570b83c60bd2de87a6fbfbebdadb4a9bd9c6cbb", size = 5350202, upload-time = "2026-08-09T13:44:58.401Z" }, + { url = "https://files.pythonhosted.org/packages/11/4b/dcd3b7eadaf4035d2c7a4289d232523a6964f602598ef7674e4bd7291f93/numpy-2.5.2-cp312-cp312-macosx_14_0_x86_64.whl", hash = "sha256:65f188481f1669e26f62b701e8205d19e460fa4a9b52a1414ba382330e4a3414", size = 6687736, upload-time = "2026-08-09T13:45:00.813Z" }, + { url = "https://files.pythonhosted.org/packages/e5/21/4947e0e9d6c9fc2e2ff15b8949049ee44f63adb9cacc729ab8793f97e712/numpy-2.5.2-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:8ee9c4eeb8454b3660a8b53493563c3e121c2fc94fbd72b848ef814ed7b676a9", size = 15612696, upload-time = "2026-08-09T13:45:04.151Z" }, + { url = "https://files.pythonhosted.org/packages/3a/5f/62d28cf019460c7f1394105b4d49d9911a9c444cb77ab0bd95a204c5a6de/numpy-2.5.2-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:3cdec01fa790a186d430433fdd4d4ffb70eed6f0eeb4bf05c8dbe2dce0a9bcb8", size = 16722264, upload-time = "2026-08-09T13:45:07.714Z" }, + { url = "https://files.pythonhosted.org/packages/14/25/3f0be4c1b9fdf5dd5e708a6806978564d7c46a055c000496309ff2a2f8af/numpy-2.5.2-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:7999d4ddb0c4025018373fd787510d46e04c769467af22869707b3c1cfd459ab", size = 16974396, upload-time = "2026-08-09T13:45:11.316Z" }, + { url = "https://files.pythonhosted.org/packages/22/72/6262cbdeeb45da9d971e40715f579d791603ba8ec0b5e2db1ac55454421d/numpy-2.5.2-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:c1f017dc0875c9209d219f97feceb7d54c2661bb243deb4114478e1295808af7", size = 18476044, upload-time = "2026-08-09T13:45:14.869Z" }, + { url = "https://files.pythonhosted.org/packages/36/33/29208b8b075bde62d26a81d14b358c42b0f69b6cabd98d4ff97f37f22b05/numpy-2.5.2-cp312-cp312-win32.whl", hash = "sha256:d6a48072864e3324e194a8fbb3c657bcc5b5c869dbc64c9537b1d5c862572c0a", size = 6072817, upload-time = "2026-08-09T13:45:17.867Z" }, + { url = "https://files.pythonhosted.org/packages/7f/b9/87fea2769fe1c47c1b5b01d8310772c9d1a85d485de7cf386ef7a3332b02/numpy-2.5.2-cp312-cp312-win_amd64.whl", hash = "sha256:28ac63476ec7651484215ee7fa15a1f78b57c14621f01e392afe17b9a1390ce4", size = 12464674, upload-time = "2026-08-09T13:45:20.734Z" }, + { url = "https://files.pythonhosted.org/packages/14/52/032b97e00461ab0809bbe4c588b035620e5a14b8cdee47ecddefc7b17d33/numpy-2.5.2-cp312-cp312-win_arm64.whl", hash = "sha256:27650bb0e7140fa3d37b9923b4803645e0b125d190f326eecfd3f4dad8e8ade1", size = 10397131, upload-time = "2026-08-09T13:45:23.73Z" }, +] + +[[package]] +name = "nvidia-cublas-cu12" +version = "12.8.4.1" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/29/99/db44d685f0e257ff0e213ade1964fc459b4a690a73293220e98feb3307cf/nvidia_cublas_cu12-12.8.4.1-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:b86f6dd8935884615a0683b663891d43781b819ac4f2ba2b0c9604676af346d0", size = 590537124, upload-time = "2025-03-07T01:43:53.556Z" }, + { url = "https://files.pythonhosted.org/packages/dc/61/e24b560ab2e2eaeb3c839129175fb330dfcfc29e5203196e5541a4c44682/nvidia_cublas_cu12-12.8.4.1-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:8ac4e771d5a348c551b2a426eda6193c19aa630236b418086020df5ba9667142", size = 594346921, upload-time = "2025-03-07T01:44:31.254Z" }, +] + +[[package]] +name = "nvidia-cuda-cupti-cu12" +version = "12.8.90" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d5/1f/b3bd73445e5cb342727fd24fe1f7b748f690b460acadc27ea22f904502c8/nvidia_cuda_cupti_cu12-12.8.90-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:4412396548808ddfed3f17a467b104ba7751e6b58678a4b840675c56d21cf7ed", size = 9533318, upload-time = "2025-03-07T01:40:10.421Z" }, + { url = "https://files.pythonhosted.org/packages/f8/02/2adcaa145158bf1a8295d83591d22e4103dbfd821bcaf6f3f53151ca4ffa/nvidia_cuda_cupti_cu12-12.8.90-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:ea0cb07ebda26bb9b29ba82cda34849e73c166c18162d3913575b0c9db9a6182", size = 10248621, upload-time = "2025-03-07T01:40:21.213Z" }, +] + +[[package]] +name = "nvidia-cuda-nvrtc-cu12" +version = "12.8.93" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/05/6b/32f747947df2da6994e999492ab306a903659555dddc0fbdeb9d71f75e52/nvidia_cuda_nvrtc_cu12-12.8.93-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:a7756528852ef889772a84c6cd89d41dfa74667e24cca16bb31f8f061e3e9994", size = 88040029, upload-time = "2025-03-07T01:42:13.562Z" }, + { url = "https://files.pythonhosted.org/packages/eb/d1/e50d0acaab360482034b84b6e27ee83c6738f7d32182b987f9c7a4e32962/nvidia_cuda_nvrtc_cu12-12.8.93-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:fc1fec1e1637854b4c0a65fb9a8346b51dd9ee69e61ebaccc82058441f15bce8", size = 43106076, upload-time = "2025-03-07T01:41:59.817Z" }, +] + +[[package]] +name = "nvidia-cuda-runtime-cu12" +version = "12.8.90" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7c/75/f865a3b236e4647605ea34cc450900854ba123834a5f1598e160b9530c3a/nvidia_cuda_runtime_cu12-12.8.90-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:52bf7bbee900262ffefe5e9d5a2a69a30d97e2bc5bb6cc866688caa976966e3d", size = 965265, upload-time = "2025-03-07T01:39:43.533Z" }, + { url = "https://files.pythonhosted.org/packages/0d/9b/a997b638fcd068ad6e4d53b8551a7d30fe8b404d6f1804abf1df69838932/nvidia_cuda_runtime_cu12-12.8.90-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:adade8dcbd0edf427b7204d480d6066d33902cab2a4707dcfc48a2d0fd44ab90", size = 954765, upload-time = "2025-03-07T01:40:01.615Z" }, +] + +[[package]] +name = "nvidia-cudnn-cu12" +version = "9.19.0.56" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-cublas-cu12" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/09/b8/277c51962ee46fa3e5b203ac5f76107c650f781d6891e681e28e6f3e9fe6/nvidia_cudnn_cu12-9.19.0.56-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:08caaf27fe556aca82a3ee3b5aa49a77e7de0cfcb7ff4e5c29da426387a8267e", size = 656910700, upload-time = "2026-02-03T20:40:25.508Z" }, + { url = "https://files.pythonhosted.org/packages/c5/41/65225d42fba06fb3dd3972485ea258e7dd07a40d6e01c95da6766ad87354/nvidia_cudnn_cu12-9.19.0.56-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:ac6ad90a075bb33a94f2b4cf4622eac13dd4dc65cf6dd9c7572a318516a36625", size = 657906812, upload-time = "2026-02-03T20:44:12.638Z" }, +] + +[[package]] +name = "nvidia-cufft-cu12" +version = "11.3.3.83" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-nvjitlink-cu12" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/60/bc/7771846d3a0272026c416fbb7e5f4c1f146d6d80704534d0b187dd6f4800/nvidia_cufft_cu12-11.3.3.83-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:848ef7224d6305cdb2a4df928759dca7b1201874787083b6e7550dd6765ce69a", size = 193109211, upload-time = "2025-03-07T01:44:56.873Z" }, + { url = "https://files.pythonhosted.org/packages/1f/13/ee4e00f30e676b66ae65b4f08cb5bcbb8392c03f54f2d5413ea99a5d1c80/nvidia_cufft_cu12-11.3.3.83-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:4d2dd21ec0b88cf61b62e6b43564355e5222e4a3fb394cac0db101f2dd0d4f74", size = 193118695, upload-time = "2025-03-07T01:45:27.821Z" }, +] + +[[package]] +name = "nvidia-cufile-cu12" +version = "1.13.1.3" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/bb/fe/1bcba1dfbfb8d01be8d93f07bfc502c93fa23afa6fd5ab3fc7c1df71038a/nvidia_cufile_cu12-1.13.1.3-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:1d069003be650e131b21c932ec3d8969c1715379251f8d23a1860554b1cb24fc", size = 1197834, upload-time = "2025-03-07T01:45:50.723Z" }, + { url = "https://files.pythonhosted.org/packages/1e/f5/5607710447a6fe9fd9b3283956fceeee8a06cda1d2f56ce31371f595db2a/nvidia_cufile_cu12-1.13.1.3-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:4beb6d4cce47c1a0f1013d72e02b0994730359e17801d395bdcbf20cfb3bb00a", size = 1120705, upload-time = "2025-03-07T01:45:41.434Z" }, +] + +[[package]] +name = "nvidia-curand-cu12" +version = "10.3.9.90" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/45/5e/92aa15eca622a388b80fbf8375d4760738df6285b1e92c43d37390a33a9a/nvidia_curand_cu12-10.3.9.90-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:dfab99248034673b779bc6decafdc3404a8a6f502462201f2f31f11354204acd", size = 63625754, upload-time = "2025-03-07T01:46:10.735Z" }, + { url = "https://files.pythonhosted.org/packages/fb/aa/6584b56dc84ebe9cf93226a5cde4d99080c8e90ab40f0c27bda7a0f29aa1/nvidia_curand_cu12-10.3.9.90-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:b32331d4f4df5d6eefa0554c565b626c7216f87a06a4f56fab27c3b68a830ec9", size = 63619976, upload-time = "2025-03-07T01:46:23.323Z" }, +] + +[[package]] +name = "nvidia-cusolver-cu12" +version = "11.7.3.90" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-cublas-cu12" }, + { name = "nvidia-cusparse-cu12" }, + { name = "nvidia-nvjitlink-cu12" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/c8/32/f7cd6ce8a7690544d084ea21c26e910a97e077c9b7f07bf5de623ee19981/nvidia_cusolver_cu12-11.7.3.90-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:db9ed69dbef9715071232caa9b69c52ac7de3a95773c2db65bdba85916e4e5c0", size = 267229841, upload-time = "2025-03-07T01:46:54.356Z" }, + { url = "https://files.pythonhosted.org/packages/85/48/9a13d2975803e8cf2777d5ed57b87a0b6ca2cc795f9a4f59796a910bfb80/nvidia_cusolver_cu12-11.7.3.90-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:4376c11ad263152bd50ea295c05370360776f8c3427b30991df774f9fb26c450", size = 267506905, upload-time = "2025-03-07T01:47:16.273Z" }, +] + +[[package]] +name = "nvidia-cusparse-cu12" +version = "12.5.8.93" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-nvjitlink-cu12" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/bc/f7/cd777c4109681367721b00a106f491e0d0d15cfa1fd59672ce580ce42a97/nvidia_cusparse_cu12-12.5.8.93-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:9b6c161cb130be1a07a27ea6923df8141f3c295852f4b260c65f18f3e0a091dc", size = 288117129, upload-time = "2025-03-07T01:47:40.407Z" }, + { url = "https://files.pythonhosted.org/packages/c2/f5/e1854cb2f2bcd4280c44736c93550cc300ff4b8c95ebe370d0aa7d2b473d/nvidia_cusparse_cu12-12.5.8.93-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:1ec05d76bbbd8b61b06a80e1eaf8cf4959c3d4ce8e711b65ebd0443bb0ebb13b", size = 288216466, upload-time = "2025-03-07T01:48:13.779Z" }, +] + +[[package]] +name = "nvidia-cusparselt-cu12" +version = "0.7.1" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/73/b9/598f6ff36faaece4b3c50d26f50e38661499ff34346f00e057760b35cc9d/nvidia_cusparselt_cu12-0.7.1-py3-none-manylinux2014_aarch64.whl", hash = "sha256:8878dce784d0fac90131b6817b607e803c36e629ba34dc5b433471382196b6a5", size = 283835557, upload-time = "2025-02-26T00:16:54.265Z" }, + { url = "https://files.pythonhosted.org/packages/56/79/12978b96bd44274fe38b5dde5cfb660b1d114f70a65ef962bcbbed99b549/nvidia_cusparselt_cu12-0.7.1-py3-none-manylinux2014_x86_64.whl", hash = "sha256:f1bb701d6b930d5a7cea44c19ceb973311500847f81b634d802b7b539dc55623", size = 287193691, upload-time = "2025-02-26T00:15:44.104Z" }, +] + +[[package]] +name = "nvidia-nccl-cu12" +version = "2.28.9" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/08/c4/120d2dfd92dff2c776d68f361ff8705fdea2ca64e20b612fab0fd3f581ac/nvidia_nccl_cu12-2.28.9-py3-none-manylinux_2_18_aarch64.whl", hash = "sha256:50a36e01c4a090b9f9c47d92cec54964de6b9fcb3362d0e19b8ffc6323c21b60", size = 296766525, upload-time = "2025-11-18T05:49:16.094Z" }, + { url = "https://files.pythonhosted.org/packages/4a/4e/44dbb46b3d1b0ec61afda8e84837870f2f9ace33c564317d59b70bc19d3e/nvidia_nccl_cu12-2.28.9-py3-none-manylinux_2_18_x86_64.whl", hash = "sha256:485776daa8447da5da39681af455aa3b2c2586ddcf4af8772495e7c532c7e5ab", size = 296782137, upload-time = "2025-11-18T05:49:34.248Z" }, +] + +[[package]] +name = "nvidia-nvjitlink-cu12" +version = "12.8.93" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f6/74/86a07f1d0f42998ca31312f998bd3b9a7eff7f52378f4f270c8679c77fb9/nvidia_nvjitlink_cu12-12.8.93-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:81ff63371a7ebd6e6451970684f916be2eab07321b73c9d244dc2b4da7f73b88", size = 39254836, upload-time = "2025-03-07T01:49:55.661Z" }, + { url = "https://files.pythonhosted.org/packages/2a/a2/8cee5da30d13430e87bf99bb33455d2724d0a4a9cb5d7926d80ccb96d008/nvidia_nvjitlink_cu12-12.8.93-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:adccd7161ace7261e01bb91e44e88da350895c270d23f744f0820c818b7229e7", size = 38386204, upload-time = "2025-03-07T01:49:43.612Z" }, +] + +[[package]] +name = "nvidia-nvshmem-cu12" +version = "3.4.5" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/1d/6a/03aa43cc9bd3ad91553a88b5f6fb25ed6a3752ae86ce2180221962bc2aa5/nvidia_nvshmem_cu12-3.4.5-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:0b48363fc6964dede448029434c6abed6c5e37f823cb43c3bcde7ecfc0457e15", size = 138936938, upload-time = "2025-09-06T00:32:05.589Z" }, + { url = "https://files.pythonhosted.org/packages/b5/09/6ea3ea725f82e1e76684f0708bbedd871fc96da89945adeba65c3835a64c/nvidia_nvshmem_cu12-3.4.5-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:042f2500f24c021db8a06c5eec2539027d57460e1c1a762055a6554f72c369bd", size = 139103095, upload-time = "2025-09-06T00:32:31.266Z" }, +] + +[[package]] +name = "nvidia-nvtx-cu12" +version = "12.8.90" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/10/c0/1b303feea90d296f6176f32a2a70b5ef230f9bdeb3a72bddb0dc922dc137/nvidia_nvtx_cu12-12.8.90-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:d7ad891da111ebafbf7e015d34879f7112832fc239ff0d7d776b6cb685274615", size = 91161, upload-time = "2025-03-07T01:42:23.922Z" }, + { url = "https://files.pythonhosted.org/packages/a2/eb/86626c1bbc2edb86323022371c39aa48df6fd8b0a1647bc274577f72e90b/nvidia_nvtx_cu12-12.8.90-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:5b17e2001cc0d751a5bc2c6ec6d26ad95913324a4adb86788c944f8ce9ba441f", size = 89954, upload-time = "2025-03-07T01:42:44.131Z" }, +] + +[[package]] +name = "packaging" +version = "26.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/7d/fa/3944b40b07da9ce895c0e6303a5ab7d53da063554f534556b134a54d6093/packaging-26.3.tar.gz", hash = "sha256:94edc256424af38762eb31306eed28beb9f0efc50a8837492c9d6fd6004aed79", size = 313412, upload-time = "2026-08-04T18:15:28.737Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/63/34/ba1c580383c9eada3711951fef0795c80b829a078d72188184bcab9dd527/packaging-26.3-py3-none-any.whl", hash = "sha256:d7193f7c8e4e93f444fde0262bf90af30e16fa0ad0ad44cb553c87339b23cd1c", size = 129956, upload-time = "2026-08-04T18:15:27.159Z" }, +] + +[[package]] +name = "pathspec" +version = "1.1.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/5a/82/42f767fc1c1143d6fd36efb827202a2d997a375e160a71eb2888a925aac1/pathspec-1.1.1.tar.gz", hash = "sha256:17db5ecd524104a120e173814c90367a96a98d07c45b2e10c2f3919fff91bf5a", size = 135180, upload-time = "2026-04-27T01:46:08.907Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f1/d9/7fb5aa316bc299258e68c73ba3bddbc499654a07f151cba08f6153988714/pathspec-1.1.1-py3-none-any.whl", hash = "sha256:a00ce642f577bf7f473932318056212bc4f8bfdf53128c78bbd5af0b9b20b189", size = 57328, upload-time = "2026-04-27T01:46:07.06Z" }, +] + +[[package]] +name = "pluggy" +version = "1.6.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f9/e2/3e91f31a7d2b083fe6ef3fa267035b518369d9511ffab804f839851d2779/pluggy-1.6.0.tar.gz", hash = "sha256:7dcc130b76258d33b90f61b658791dede3486c3e6bfb003ee5c9bfb396dd22f3", size = 69412, upload-time = "2025-05-15T12:30:07.975Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl", hash = "sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746", size = 20538, upload-time = "2025-05-15T12:30:06.134Z" }, +] + +[[package]] +name = "pytest" +version = "8.3.5" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "iniconfig" }, + { name = "packaging" }, + { name = "pluggy" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/ae/3c/c9d525a414d506893f0cd8a8d0de7706446213181570cdbd766691164e40/pytest-8.3.5.tar.gz", hash = "sha256:f4efe70cc14e511565ac476b57c279e12a855b11f48f212af1080ef2263d3845", size = 1450891, upload-time = "2025-03-02T12:54:54.503Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/30/3d/64ad57c803f1fa1e963a7946b6e0fea4a70df53c1a7fed304586539c2bac/pytest-8.3.5-py3-none-any.whl", hash = "sha256:c69214aa47deac29fad6c2a4f590b9c4a9fdb16a403176fe154b79c0b4d4d820", size = 343634, upload-time = "2025-03-02T12:54:52.069Z" }, +] + +[[package]] +name = "pytest-mock" +version = "3.15.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pytest" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/68/14/eb014d26be205d38ad5ad20d9a80f7d201472e08167f0bb4361e251084a9/pytest_mock-3.15.1.tar.gz", hash = "sha256:1849a238f6f396da19762269de72cb1814ab44416fa73a8686deac10b0d87a0f", size = 34036, upload-time = "2025-09-16T16:37:27.081Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5a/cc/06253936f4a7fa2e0f48dfe6d851d9c56df896a9ab09ac019d70b760619c/pytest_mock-3.15.1-py3-none-any.whl", hash = "sha256:0a25e2eb88fe5168d535041d09a4529a188176ae608a6d249ee65abc0949630d", size = 10095, upload-time = "2025-09-16T16:37:25.734Z" }, +] + +[[package]] +name = "regex" +version = "2026.9.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/19/c1/6b30b775c7bcc6cf6506a4d4741c2123e8d99cd50f3fe8cbd731f5fef526/regex-2026.9.3.tar.gz", hash = "sha256:aabd43208e335f4c3f0b56de3464b066dd425983a58f6eeb5738bcd7465403db", size = 416720, upload-time = "2026-09-01T00:53:43.821Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/da/cb/cba530bc3b068fc337f8f455c63ef5ee91a4eb4c76ecf5998e5cef5aaa6b/regex-2026.9.3-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:5db80d0b1c8238940b5957dd66b5c818ea40a221f6652fb717c027a562d09c77", size = 496699, upload-time = "2026-09-01T00:50:27.98Z" }, + { url = "https://files.pythonhosted.org/packages/81/39/f2e9fb6bbbc80f8bf67ad79d7e2e8866f7837d7c24c692f7faf8f1272e7e/regex-2026.9.3-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:35d48ce3dee087b63b15cd0a7a3110d0a76c29edbe1f2ad0520b8c4adb7cb596", size = 297018, upload-time = "2026-09-01T00:50:29.487Z" }, + { url = "https://files.pythonhosted.org/packages/a1/b6/c16ee58840baf7659def27ef6f62f3d9a9909670d3c1b4b98bb8b8ee47e2/regex-2026.9.3-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:1f22e0d21ae7016c77175c139a7fca465b988efc1280df4816c79752068d9e2e", size = 292008, upload-time = "2026-09-01T00:50:30.929Z" }, + { url = "https://files.pythonhosted.org/packages/cb/b4/4987bf0f17604669b4ea5aef219886d0a73188c4716ff3a7d275d4d15c15/regex-2026.9.3-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:233662cf8cfdfe3c0e58aa8f7bbefc579b5be0ac34546f123c159804179e8687", size = 796101, upload-time = "2026-09-01T00:50:32.486Z" }, + { url = "https://files.pythonhosted.org/packages/0b/95/2a9ab02a68c8a61dc0b4882ed643b1a95740d9dc291dc26c77d19af79691/regex-2026.9.3-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:7d2eed2e4d231278a2ccab3f4bfa2c1e39855f336475f7756a281d767d2b1753", size = 865435, upload-time = "2026-09-01T00:50:34.171Z" }, + { url = "https://files.pythonhosted.org/packages/04/92/0570d41559b446c97c1148cb9ebc1df09f2949b03c7c9bfee09976b3465f/regex-2026.9.3-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:5e674cecb61cb160be392da07fd8a71509ef927f437fbf3215432692ed385151", size = 911828, upload-time = "2026-09-01T00:50:35.72Z" }, + { url = "https://files.pythonhosted.org/packages/4b/8b/9cc6d4123033f7cb82df6cd8ce19eb0fc18a964afe060a03c9b26757c9f3/regex-2026.9.3-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:665207e41bacd435db001099eeab44103197c2c1a729d73ade74688a905ed4ce", size = 801965, upload-time = "2026-09-01T00:50:37.701Z" }, + { url = "https://files.pythonhosted.org/packages/c9/98/39262e91aa87a67c82cbe90a0df4c3d382c7a44811fe80067904085211b4/regex-2026.9.3-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:7a7ddc9a8ca1795166a1ca80364b8ce74187fc210e112d3fb048b711b934f36c", size = 776192, upload-time = "2026-09-01T00:50:39.57Z" }, + { url = "https://files.pythonhosted.org/packages/24/e9/3bb93fe4ee4b6f8ce7ba69b527c4a63cfa3393fc425ab26486041fe441c8/regex-2026.9.3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:e3037d02425863ce9501afbaa04ba967162810004bacde39a53ea9a5b740eb32", size = 785053, upload-time = "2026-09-01T00:50:41.156Z" }, + { url = "https://files.pythonhosted.org/packages/f9/05/31d5bc2553a700c0dfc6b5b6a13c61cdcd1210fde1e304cfa18a33f138b2/regex-2026.9.3-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:3de4eab8c763393b75bbb26f81934ab2cc8794f48f79e90622e3ab7ea57f3d14", size = 860546, upload-time = "2026-09-01T00:50:42.746Z" }, + { url = "https://files.pythonhosted.org/packages/65/a3/2e1e854d80becda0f061093805bbfc037a5849448f46d0a2b71a070d45e2/regex-2026.9.3-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:98620c9c4c22568ad70f57b80527c780b6f8fd26e36507bf8e2273262a228275", size = 765841, upload-time = "2026-09-01T00:50:44.5Z" }, + { url = "https://files.pythonhosted.org/packages/6a/d6/43d02948cedde2e8476ac893ea02755ee5ee1b21c531fda92d80e114f0bc/regex-2026.9.3-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:0b1ba3aaaf5776de473ee16625ac60ac195abb0343afb273575a8201d99be089", size = 852147, upload-time = "2026-09-01T00:50:46.474Z" }, + { url = "https://files.pythonhosted.org/packages/21/ff/adb4e2d08afe8f4c6df004d94604257e1f72af7ba328af7715601585aba4/regex-2026.9.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:56d8659c65166641d8f1b5efccc391c62c8a899eff4d528b981cc62b7b402a4b", size = 789761, upload-time = "2026-09-01T00:50:48.749Z" }, + { url = "https://files.pythonhosted.org/packages/c3/e1/1490d1351758e87f6e702cf2025036bdc7bc59182e2ff5c7bec004b19aed/regex-2026.9.3-cp312-cp312-win32.whl", hash = "sha256:837c1859913798d8bebcd98d4a037e113f8d79e81733009bf590e449769eecb3", size = 267150, upload-time = "2026-09-01T00:50:50.414Z" }, + { url = "https://files.pythonhosted.org/packages/d5/49/4c40cf722d84d60e807a08ef4c3f579216bf97df60c4a1b10be49655d302/regex-2026.9.3-cp312-cp312-win_amd64.whl", hash = "sha256:1ba1dbbb93c5c5629c1861763aec5bfa9f05ad24ef450694130e25029ce7bc36", size = 277773, upload-time = "2026-09-01T00:50:51.963Z" }, + { url = "https://files.pythonhosted.org/packages/aa/af/c48b3b2b4244b4b090554c78d3387e9ae7b859f3dbf7148a27d427e9e5b8/regex-2026.9.3-cp312-cp312-win_arm64.whl", hash = "sha256:d7b3a8a4bbd83ad8b29758f5d24bab10a3f2de87970db36f1e3651c733353136", size = 277122, upload-time = "2026-09-01T00:50:53.778Z" }, +] + +[[package]] +name = "setuptools" +version = "81.0.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/0d/1c/73e719955c59b8e424d015ab450f51c0af856ae46ea2da83eba51cc88de1/setuptools-81.0.0.tar.gz", hash = "sha256:487b53915f52501f0a79ccfd0c02c165ffe06631443a886740b91af4b7a5845a", size = 1198299, upload-time = "2026-02-06T21:10:39.601Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e1/e3/c164c88b2e5ce7b24d667b9bd83589cf4f3520d97cad01534cd3c4f55fdb/setuptools-81.0.0-py3-none-any.whl", hash = "sha256:fdd925d5c5d9f62e4b74b30d6dd7828ce236fd6ed998a08d81de62ce5a6310d6", size = 1062021, upload-time = "2026-02-06T21:10:37.175Z" }, +] + +[[package]] +name = "sympy" +version = "1.14.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "mpmath" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/83/d3/803453b36afefb7c2bb238361cd4ae6125a569b4db67cd9e79846ba2d68c/sympy-1.14.0.tar.gz", hash = "sha256:d3d3fe8df1e5a0b42f0e7bdf50541697dbe7d23746e894990c030e2b05e72517", size = 7793921, upload-time = "2025-04-27T18:05:01.611Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a2/09/77d55d46fd61b4a135c444fc97158ef34a095e5681d0a6c10b75bf356191/sympy-1.14.0-py3-none-any.whl", hash = "sha256:e091cc3e99d2141a0ba2847328f5479b05d94a6635cb96148ccb3f34671bd8f5", size = 6299353, upload-time = "2025-04-27T18:04:59.103Z" }, +] + +[[package]] +name = "torch" +version = "2.11.0+cu128" +source = { registry = "https://download.pytorch.org/whl/cu128" } +dependencies = [ + { name = "cuda-bindings", marker = "sys_platform == 'linux'" }, + { name = "cuda-toolkit", extra = ["cublas", "cudart", "cufft", "cufile", "cupti", "curand", "cusolver", "cusparse", "nvjitlink", "nvrtc", "nvtx"], marker = "sys_platform == 'linux'" }, + { name = "filelock" }, + { name = "fsspec" }, + { name = "jinja2" }, + { name = "networkx" }, + { name = "nvidia-cudnn-cu12", marker = "sys_platform == 'linux'" }, + { name = "nvidia-cusparselt-cu12", marker = "sys_platform == 'linux'" }, + { name = "nvidia-nccl-cu12", marker = "sys_platform == 'linux'" }, + { name = "nvidia-nvshmem-cu12", marker = "sys_platform == 'linux'" }, + { name = "setuptools" }, + { name = "sympy" }, + { name = "triton", marker = "sys_platform == 'linux'" }, + { name = "typing-extensions" }, +] +wheels = [ + { url = "https://download-r2.pytorch.org/whl/cu128/torch-2.11.0%2Bcu128-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:9c8f38efee365cb9d334de8a83ce52fc7e5fc9e5a7b0853285efa1b69e00b0f2", upload-time = "2026-04-27T17:41:30Z" }, + { url = "https://download-r2.pytorch.org/whl/cu128/torch-2.11.0%2Bcu128-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:d252cf975fb18c94a85336323ad425f473df56dab35a44b00399bd70c7a3b997", upload-time = "2026-04-27T17:42:06Z" }, + { url = "https://download-r2.pytorch.org/whl/cu128/torch-2.11.0%2Bcu128-cp312-cp312-win_amd64.whl", hash = "sha256:7c78215c3af4f62e63f2b2e360f1722fc719b0853c7ac22666483d9810613a4c", upload-time = "2026-04-27T17:43:49Z" }, +] + +[[package]] +name = "triton" +version = "3.6.0" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/17/5d/08201db32823bdf77a0e2b9039540080b2e5c23a20706ddba942924ebcd6/triton-3.6.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:374f52c11a711fd062b4bfbb201fd9ac0a5febd28a96fb41b4a0f51dde3157f4", size = 176128243, upload-time = "2026-01-20T16:16:07.857Z" }, + { url = "https://files.pythonhosted.org/packages/ab/a8/cdf8b3e4c98132f965f88c2313a4b493266832ad47fb52f23d14d4f86bb5/triton-3.6.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:74caf5e34b66d9f3a429af689c1c7128daba1d8208df60e81106b115c00d6fca", size = 188266850, upload-time = "2026-01-20T16:00:43.041Z" }, +] + +[[package]] +name = "typing-extensions" +version = "4.16.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f6/cc/6253133b5bb138fc3306cebfbda2c520f545d36b5be2c7255cc528bb45d6/typing_extensions-4.16.0.tar.gz", hash = "sha256:dc983d19a509c94dba722ee6abd33940f7c05a89e243c47e907eb4db6f1a43e5", size = 113555, upload-time = "2026-07-02T08:40:05.92Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/49/d3/b8441a820a491ddfc024b0b0cf0393375b75ea13866d9c66727e54c2fc80/typing_extensions-4.16.0-py3-none-any.whl", hash = "sha256:481caa481374e813c1b176ada14e97f1f67a4539ce9cfeb3f350d78d6370c2e8", size = 45571, upload-time = "2026-07-02T08:40:04.659Z" }, +] + +[[package]] +name = "wcwidth" +version = "0.8.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/36/57/ed58088fafdf4c55a0ad6bde846502567645424d7ebf325230b9237f4085/wcwidth-0.8.3.tar.gz", hash = "sha256:d128512515fbf4612e0ff21fd6380399210318b7b54a9af59dff8454cf9730eb", size = 1458450, upload-time = "2026-08-28T18:10:06.875Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c4/0e/57f6bb3024a597b2e8ec4aee710ffe62ddc95af2e2bb1ee7a7abdc22c68c/wcwidth-0.8.3-py3-none-any.whl", hash = "sha256:d5b73dba6158a595ec9370350e7f2637bcac8d6c5e4fde34f30fcffb6103a5e4", size = 331669, upload-time = "2026-08-28T18:10:04.909Z" }, +] diff --git a/dimos/imitation/policy/abc/test_module.py b/dimos/imitation/policy/abc/test_module.py new file mode 100644 index 0000000000..b98df335b3 --- /dev/null +++ b/dimos/imitation/policy/abc/test_module.py @@ -0,0 +1,45 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from pathlib import Path + +from dimos.experimental.isolated_python.module import contract_rpc_names +from dimos.imitation.policy.abc.module import DualOpenYamAbcPolicy + + +def test_abc_contract_is_importable_without_torch() -> None: + blueprint = DualOpenYamAbcPolicy.blueprint(artifact="checkpoint.pt", task="bottles") + streams = {stream.name for stream in blueprint.blueprints[0].streams} + + assert streams == { + "button_pressed", + "top_image", + "left_wrist_image", + "right_wrist_image", + "coordinator_joint_state", + } + assert contract_rpc_names(DualOpenYamAbcPolicy) == { + "preflight_rollout", + "rollout_status", + "start_rollout", + "stop_rollout", + } + + +def test_abc_contract_resolves_its_own_isolated_project() -> None: + module = DualOpenYamAbcPolicy(artifact="checkpoint.pt", task="bottles") + try: + assert module.runtime_project == Path(__file__).parent / "python" + finally: + module.stop() diff --git a/dimos/imitation/policy/backend.py b/dimos/imitation/policy/backend.py new file mode 100644 index 0000000000..d0db762265 --- /dev/null +++ b/dimos/imitation/policy/backend.py @@ -0,0 +1,53 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Small in-process contract implemented by isolated policy backends.""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Any, Protocol + +import numpy as np +from numpy.typing import NDArray + +from dimos.imitation.profile import PolicyIOProfile + + +@dataclass(frozen=True) +class PolicyBackendInfo: + """Execution information discovered while loading a policy artifact.""" + + name: str + chunk_length: int + preferred_execution_steps: int + action_lower: NDArray[np.float32] | None = None + action_upper: NDArray[np.float32] | None = None + + +class PolicyBackend(Protocol): + """Backend-specific loading and inference behind the common rollout loop.""" + + def __init__(self, config: Any) -> None: ... + + def load(self, profile: PolicyIOProfile) -> PolicyBackendInfo: ... + + def reset(self) -> None: ... + + def predict( + self, + observations: Mapping[str, NDArray[Any]], + task: str, + ) -> NDArray[np.float32]: ... diff --git a/dimos/imitation/policy/lerobot/README.md b/dimos/imitation/policy/lerobot/README.md index d53fc8f3e1..de40e22aef 100644 --- a/dimos/imitation/policy/lerobot/README.md +++ b/dimos/imitation/policy/lerobot/README.md @@ -1,61 +1,26 @@ -# LeRobot Policy Module +# LeRobot Policy Backend -`LeRobotPolicyModule` runs trained LeRobot policies in a managed Python-native -subprocess. Its LeRobot, Transformers, Torch, and NumPy versions live in the -sibling `python/` project and do not change the main DimOS environment. - -The host contract subscribes to: - -- `color_image: Image` -- `coordinator_joint_state: JointState` -- `button_pressed: Buttons` - -It submits complete, timestamped action chunks to one named -`JointTrajectoryTask` through the control coordinator. The state and action -vectors use `joint_names` order, including any gripper joint. The policy output -is already postprocessed into each joint's native absolute coordinate; the -runtime does not reinterpret gripper values. +`OpenYamLeRobotPolicy` is generated from `OPENYAM_QUEST_IO`. Its host process +contains no LeRobot imports; the sibling locked project implements the common +policy backend and runs prediction in the same isolated process as the shared +rollout loop. ```python -from dimos.imitation.policy.lerobot.module import LeRobotPolicyModule +from dimos.imitation.policy.lerobot.module import OpenYamLeRobotPolicy -policy = LeRobotPolicyModule.blueprint( - policy_path="outputs/pick/checkpoints/last/pretrained_model", +policy = OpenYamLeRobotPolicy.blueprint( + instance_name="PolicyRolloutModule", + artifact="outputs/pick/checkpoints/last/pretrained_model", task="pick up the object", - joint_names=["arm/joint1", "arm/joint2", "arm/gripper"], - trajectory_task_name="policy_rollout", - fps=30.0, - robot_type="my_robot", - image_width=640, - image_height=480, ) ``` -The module exposes `preflight_rollout`, `start_rollout`, `stop_rollout`, and -`rollout_status` RPCs. Preflight loads the checkpoint and processors, validates -the control task and fresh live observations, and sends no trajectory. -`start_rollout` refuses to run until preflight passes and rechecks observations -before starting. -The runtime rejects missing or stale observations, missing joints, non-finite -values, incompatible checkpoint features, malformed action chunks, and -trajectories outside the hardware's declared position limits. Pressing the -configured Quest button (A by default) toggles a preflighted rollout. - -The runtime calls LeRobot's `predict_action_chunk()`, postprocesses the entire -chunk, clips every action dimension to the checkpoint's recorded data range, -and executes its first `n_action_steps` at the configured `fps`. The coordinator -still validates the resulting trajectory against hardware limits. Each trajectory -starts with the joint-state observation used for inference, so the coordinator -rejects a stale start if the robot moved in the meantime. Configure `fps` to -match the action frequency used by the training dataset. - -Current limitation: this contract assumes every postprocessed action is an -absolute target in the connected hardware joint's native coordinate. A generic -contract for checkpoints that encode grippers in normalized or device-specific -coordinates remains future work; this runtime does not special-case those -grippers. +The profile supplies `wrist_image`, `coordinator_joint_state`, feature keys, +image shape, joint order, and 30 Hz rate. The LeRobot adapter validates those +keys and dimensions, loads pre/postprocessors, returns a 2-D action chunk, and +reports checkpoint action bounds to the common safety loop. -Run isolated runtime checks with: +Run isolated checks with: ```bash cd dimos/imitation/policy/lerobot/python diff --git a/dimos/imitation/policy/lerobot/python/dimos_lerobot/runtime.py b/dimos/imitation/policy/lerobot/python/dimos_lerobot/runtime.py index 9cd5ce5c48..4f61a78d14 100644 --- a/dimos/imitation/policy/lerobot/python/dimos_lerobot/runtime.py +++ b/dimos/imitation/policy/lerobot/python/dimos_lerobot/runtime.py @@ -12,14 +12,13 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Run trained LeRobot policies in an isolated Python environment.""" +"""LeRobot adapter for the shared DimOS policy rollout runtime.""" from __future__ import annotations +from collections.abc import Mapping from contextlib import nullcontext from dataclasses import dataclass -from threading import Condition, Event, RLock, Thread, current_thread -import time from typing import Any from lerobot.configs.policies import PreTrainedConfig @@ -31,30 +30,15 @@ from lerobot.utils.import_utils import register_third_party_plugins import numpy as np from numpy.typing import NDArray -from reactivex.disposable import Disposable import torch -from dimos.constants import DEFAULT_THREAD_JOIN_TIMEOUT -from dimos.control.tasks.trajectory_task.trajectory_task import TrajectoryExecutionStatus -from dimos.core.core import rpc +from dimos.imitation.policy.backend import PolicyBackendInfo from dimos.imitation.policy.lerobot.module import ( - LeRobotPolicyModule, - RolloutStatus, + LeRobotPolicyConfig, + OpenYamLeRobotPolicy, ) -from dimos.msgs.sensor_msgs.Image import Image, ImageFormat -from dimos.msgs.sensor_msgs.JointState import JointState -from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory -from dimos.msgs.trajectory_msgs.TrajectoryPoint import TrajectoryPoint -from dimos.teleop.quest.quest_types import BUTTON_ALIASES, Buttons -from dimos.utils.logging_config import setup_logger - -logger = setup_logger() - -_IMAGE_FEATURE = "observation.images.wrist" -_STATE_FEATURE = "observation.state" -_ACTION_FEATURE = "action" - -RawObservation = dict[str, NDArray[np.uint8] | NDArray[np.float32]] +from dimos.imitation.policy.runtime import declare_policy_runtime +from dimos.imitation.profile import ImageSource, PolicyIOProfile @dataclass(frozen=True) @@ -64,472 +48,130 @@ class _LoadedPolicy: preprocessor: PolicyProcessorPipeline[RobotObservation, RobotObservation] postprocessor: PolicyProcessorPipeline[PolicyAction, PolicyAction] use_amp: bool - chunk_size: int | None - n_action_steps: int - action_lower: NDArray[np.float32] - action_upper: NDArray[np.float32] - - -class LeRobotPolicyRuntime(LeRobotPolicyModule): - """Concrete LeRobot implementation loaded by ``LeRobotPolicyModule``.""" - - _lock: RLock - _observation_changed: Condition - _loaded_policy: _LoadedPolicy | None - _latest_image: tuple[NDArray[np.uint8], float] | None - _latest_joint_state: JointState | None - _stop_event: Event - _thread: Thread | None - _chunks_accepted: int - _last_error: str | None - _active: bool - - def __init__(self, **kwargs: Any) -> None: - super().__init__(**kwargs) - self._lock = RLock() - self._observation_changed = Condition(self._lock) - self._loaded_policy = None - self._latest_image = None - self._latest_joint_state = None - self._stop_event = Event() - self._thread = None - self._chunks_accepted = 0 - self._last_error = None - self._active = False - - @rpc - def start(self) -> None: - super().start() - self.register_disposable(Disposable(self.color_image.subscribe(self._on_color_image))) - self.register_disposable( - Disposable(self.coordinator_joint_state.subscribe(self._on_joint_state)) - ) - self.register_disposable(Disposable(self.button_pressed.subscribe(self._on_button_pressed))) - - @rpc - def stop(self) -> None: - if not self._stop_policy(): - self._cancel_after_stop_timeout() - super().stop() - - @rpc - def preflight_rollout(self) -> RolloutStatus: - """Validate the checkpoint, coordinator, and live observations without moving.""" - with self._lock: - if self._active: - self._last_error = "cannot preflight while a policy rollout is active" - return self._status_locked() - loaded_policy = self._loaded_policy - try: - self._snapshot_observation(time.time()) - except Exception as exc: - self._loaded_policy = None - self._last_error = str(exc) - return self._status_locked() - - try: - tasks = set(self._control.list_tasks()) - if self.config.trajectory_task_name not in tasks: - raise RuntimeError( - "ControlCoordinator is missing configured rollout task " - f"{self.config.trajectory_task_name!r}" - ) - if loaded_policy is None: - loaded_policy = self._load_policy() - logger.info( - "Loaded LeRobot policy during preflight", - path=self.config.policy_path, - runtime_fps=self.config.fps, - chunk_size=loaded_policy.chunk_size, - n_action_steps=loaded_policy.n_action_steps, - ) - with self._lock: - self._loaded_policy = loaded_policy - self._snapshot_observation(time.time()) - self._last_error = None - return self._status_locked() - except Exception as exc: - with self._lock: - self._loaded_policy = None - self._last_error = str(exc) - return self._status_locked() - @rpc - def start_rollout(self) -> RolloutStatus: - with self._lock: - if self._thread is not None and self._thread.is_alive(): - self._last_error = "a policy rollout is already active" - return self._status_locked() - if self._loaded_policy is None: - self._last_error = "policy preflight has not passed" - return self._status_locked() - try: - self._snapshot_observation(time.time()) - except RuntimeError as exc: - self._last_error = str(exc) - return self._status_locked() - self._stop_event.clear() - self._chunks_accepted = 0 - self._last_error = None - self._active = True - self._thread = Thread( - target=self._run_rollout, - name="lerobot-policy-rollout", - daemon=True, - ) - self._thread.start() - return self._status_locked() - - @rpc - def stop_rollout(self) -> RolloutStatus: - if not self._stop_policy(): - self._cancel_after_stop_timeout() - return self.rollout_status() - - @rpc - def rollout_status(self) -> RolloutStatus: - with self._lock: - return self._status_locked() - - def _status_locked(self) -> RolloutStatus: - try: - self._snapshot_observation(time.time()) - observations_ready = True - except RuntimeError: - observations_ready = False - return { - "active": self._active, - "policy_path": self.config.policy_path, - "task": self.config.task, - "device": self.config.device, - "policy_ready": self._loaded_policy is not None, - "observations_ready": observations_ready, - "chunks_accepted": self._chunks_accepted, - "last_error": self._last_error, - } - - def _on_color_image(self, image: Image) -> None: - if image.format != ImageFormat.RGB or image.data.dtype != np.uint8: - logger.warning("Ignoring non-uint8 RGB policy image", image=str(image)) - return - expected_shape = (self.config.image_height, self.config.image_width, 3) - if image.data.shape != expected_shape: - logger.warning( - "Ignoring policy image with unexpected shape", - shape=image.data.shape, - expected=expected_shape, - ) - return - with self._lock: - self._latest_image = (np.ascontiguousarray(image.data), image.ts) +class LeRobotBackend: + """Translate the generic profile-keyed arrays to a LeRobot checkpoint.""" - def _on_joint_state(self, state: JointState) -> None: - with self._observation_changed: - self._latest_joint_state = JointState(state) - self._observation_changed.notify_all() + def __init__(self, config: LeRobotPolicyConfig) -> None: + self.config = config + self._loaded: _LoadedPolicy | None = None + self._profile: PolicyIOProfile | None = None - def _on_button_pressed(self, buttons: Buttons) -> None: - button = BUTTON_ALIASES.get(self.config.rollout_button, self.config.rollout_button) - if not bool(getattr(buttons, button)): - return - with self._lock: - active = self._active - if active: - self.stop_rollout() - else: - self.start_rollout() - - def _snapshot_observation( - self, now: float - ) -> tuple[NDArray[np.uint8], NDArray[np.float32], float]: - if self._latest_image is None: - raise RuntimeError("no camera image has been received") - if self._latest_joint_state is None: - raise RuntimeError("no coordinator joint state has been received") - - image, image_ts = self._latest_image - state = self._latest_joint_state - max_age = self.config.max_observation_age_s - if now - image_ts > max_age: - raise RuntimeError(f"camera image is stale by {now - image_ts:.2f}s") - if now - state.ts > max_age: - raise RuntimeError(f"joint state is stale by {now - state.ts:.2f}s") - - positions = dict(zip(state.name, state.position, strict=False)) - missing = [name for name in self.config.joint_names if name not in positions] - if missing: - raise RuntimeError(f"joint state is missing configured joints: {missing}") - vector = np.asarray( - [positions[name] for name in self.config.joint_names], - dtype=np.float32, - ) - if not np.all(np.isfinite(vector)): - raise RuntimeError("joint state contains non-finite positions") - return image.copy(), vector, state.ts - - def _load_policy(self) -> _LoadedPolicy: + def load(self, profile: PolicyIOProfile) -> PolicyBackendInfo: register_third_party_plugins() - policy_config = PreTrainedConfig.from_pretrained(self.config.policy_path) + policy_config = PreTrainedConfig.from_pretrained(self.config.artifact) if self.config.device is not None: policy_config.device = self.config.device if policy_config.device is None: raise RuntimeError("LeRobot did not resolve an inference device") + _validate_features(policy_config, profile) - self._validate_features(policy_config) device = torch.device(policy_config.device) if device.type == "cuda" and not torch.cuda.is_available(): raise RuntimeError( f"Policy requested device {policy_config.device!r}, but CUDA is not available" ) - policy_class = get_policy_class(policy_config.type) - loaded_policy = policy_class.from_pretrained(self.config.policy_path, config=policy_config) + policy = policy_class.from_pretrained(self.config.artifact, config=policy_config) preprocessor, postprocessor = make_pre_post_processors( policy_cfg=policy_config, - pretrained_path=self.config.policy_path, + pretrained_path=self.config.artifact, preprocessor_overrides={"device_processor": {"device": str(device)}}, ) - action_lower, action_upper = _checkpoint_action_bounds( - postprocessor, - len(self.config.joint_names), - ) - return _LoadedPolicy( - policy=loaded_policy, + width = len(profile.action.demonstration.joints) + lower, upper = _checkpoint_action_bounds(postprocessor, width) + n_action_steps = _positive_int_attribute(policy_config, "n_action_steps") + chunk_length = _optional_int_attribute(policy_config, "chunk_size") or n_action_steps + self._loaded = _LoadedPolicy( + policy=policy, device=device, preprocessor=preprocessor, postprocessor=postprocessor, use_amp=bool(policy_config.use_amp), - chunk_size=_optional_int_attribute(policy_config, "chunk_size"), - n_action_steps=_positive_int_attribute(policy_config, "n_action_steps"), - action_lower=action_lower, - action_upper=action_upper, + ) + self._profile = profile + return PolicyBackendInfo( + name="lerobot", + chunk_length=chunk_length, + preferred_execution_steps=n_action_steps, + action_lower=lower, + action_upper=upper, ) - def _validate_features(self, policy_config: PreTrainedConfig) -> None: - inputs = policy_config.input_features or {} - outputs = policy_config.output_features or {} - missing = {_IMAGE_FEATURE, _STATE_FEATURE} - set(inputs) - if missing: - raise ValueError( - "Policy is incompatible with the DimOS single-camera runtime; " - f"missing input features: {sorted(missing)}" - ) - if _ACTION_FEATURE not in outputs: - raise ValueError(f"Policy has no {_ACTION_FEATURE!r} output feature") - if getattr(policy_config, "temporal_ensemble_coeff", None) is not None: - raise ValueError("Policies using temporal ensembling are not supported") - - state_shape = tuple(inputs[_STATE_FEATURE].shape) - image_shape = tuple(inputs[_IMAGE_FEATURE].shape) - action_shape = tuple(outputs[_ACTION_FEATURE].shape) - joint_count = len(self.config.joint_names) - expected_image_shape = (3, self.config.image_height, self.config.image_width) - if image_shape != expected_image_shape: - raise ValueError( - f"Policy image shape {image_shape} does not match {expected_image_shape}" - ) - if not state_shape or state_shape[0] != joint_count: - raise ValueError( - f"Policy state dimension {state_shape} does not match {joint_count} configured joints" - ) - if not action_shape or action_shape[0] != joint_count: - raise ValueError( - f"Policy action dimension {action_shape} does not match {joint_count} configured joints" - ) + def reset(self) -> None: + loaded = self._require_loaded() + _reset(loaded.policy) + _reset(loaded.preprocessor) + _reset(loaded.postprocessor) - def _predict( + def predict( self, - loaded_policy: _LoadedPolicy, - image: NDArray[np.uint8], - state: NDArray[np.float32], - *, + observations: Mapping[str, NDArray[Any]], task: str, ) -> NDArray[np.float32]: - observation: RawObservation = { - _IMAGE_FEATURE: image, - _STATE_FEATURE: state, - } + loaded = self._require_loaded() + assert self._profile is not None with ( torch.inference_mode(), torch.autocast(device_type="cuda") - if loaded_policy.device.type == "cuda" and loaded_policy.use_amp + if loaded.device.type == "cuda" and loaded.use_amp else nullcontext(), ): prepared = prepare_observation_for_inference( - observation, - loaded_policy.device, + dict(observations), + loaded.device, task=task, - robot_type=self.config.robot_type, + robot_type=self._profile.robot_type, ) - prepared = loaded_policy.preprocessor(prepared) - predict = getattr(loaded_policy.policy, "predict_action_chunk", None) + prepared = loaded.preprocessor(prepared) + predict = getattr(loaded.policy, "predict_action_chunk", None) if not callable(predict): raise TypeError("Policy does not provide predict_action_chunk()") - action_chunk = loaded_policy.postprocessor(predict(prepared)) - return np.asarray(action_chunk.to("cpu").numpy(), dtype=np.float32) - - def _run_rollout(self) -> None: - loaded_policy: _LoadedPolicy | None = None - try: - with self._lock: - loaded_policy = self._loaded_policy - if loaded_policy is None: - raise RuntimeError("policy preflight has not passed") - if self._stop_event.is_set(): - return - self._reset_policy(loaded_policy) - - while not self._stop_event.is_set(): - with self._lock: - image, state, state_ts = self._snapshot_observation(time.time()) - action_chunk = self._predict(loaded_policy, image, state, task=self.config.task) - expected_width = len(self.config.joint_names) - if action_chunk.ndim != 3 or action_chunk.shape[0] != 1: - raise RuntimeError( - f"policy returned action chunk shape {action_chunk.shape}, expected " - f"(1, steps, {expected_width})" - ) - if action_chunk.shape[2] != expected_width: - raise RuntimeError( - f"policy returned action width {action_chunk.shape[2]}, expected {expected_width}" - ) - if action_chunk.shape[1] < loaded_policy.n_action_steps: - raise RuntimeError( - f"policy returned {action_chunk.shape[1]} action steps, but n_action_steps " - f"is {loaded_policy.n_action_steps}" - ) - actions = action_chunk[0, : loaded_policy.n_action_steps] - if not np.all(np.isfinite(actions)): - raise RuntimeError("policy returned non-finite joint targets") - bounded_actions = np.clip( - actions, - loaded_policy.action_lower, - loaded_policy.action_upper, - ) - clipped = np.any(actions != bounded_actions, axis=0) - if np.any(clipped): - logger.warning( - "Clipped policy actions to checkpoint range", - joints=[ - name - for name, was_clipped in zip( - self.config.joint_names, clipped, strict=True - ) - if was_clipped - ], - ) - actions = bounded_actions - if self._stop_event.is_set(): - break - result = self._control.execute_trajectory( - self._trajectory(state, actions), - task_name=self.config.trajectory_task_name, - ) - if result.status is TrajectoryExecutionStatus.START_STATE_MISMATCH: - self._wait_for_newer_joint_state(state_ts) - continue - if result.status is not TrajectoryExecutionStatus.ACCEPTED: - raise RuntimeError( - result.message or f"trajectory rejected: {result.status.name}" - ) - with self._lock: - self._chunks_accepted += 1 - self._stop_event.wait(loaded_policy.n_action_steps / self.config.fps) - except Exception as exc: - with self._lock: - self._last_error = str(exc) - logger.exception("LeRobot policy execution stopped", error=str(exc)) - finally: - self._stop_event.set() - cancellation_error = self._cancel_trajectory() - if loaded_policy is not None: - self._reset_policy(loaded_policy) - with self._lock: - if cancellation_error is not None: - self._last_error = ( - f"{self._last_error}; {cancellation_error}" - if self._last_error is not None - else cancellation_error - ) - self._active = False - - @staticmethod - def _reset_policy(loaded_policy: _LoadedPolicy) -> None: - _reset(loaded_policy.policy) - _reset(loaded_policy.preprocessor) - _reset(loaded_policy.postprocessor) - - def _stop_policy(self) -> bool: - with self._lock: - thread = self._thread - self._stop_event.set() - self._observation_changed.notify_all() - if thread is not None and thread is not current_thread(): - thread.join(timeout=DEFAULT_THREAD_JOIN_TIMEOUT) - return thread is None or not thread.is_alive() - - def _cancel_after_stop_timeout(self) -> None: - timeout_error = f"policy rollout did not stop within {DEFAULT_THREAD_JOIN_TIMEOUT} seconds" - cancellation_error = self._cancel_trajectory() - with self._lock: - self._last_error = ( - f"{timeout_error}; {cancellation_error}" - if cancellation_error is not None - else timeout_error - ) - - def _trajectory( - self, - state: NDArray[np.float32], - actions: NDArray[np.float32], - ) -> JointTrajectory: - zeros = [0.0] * len(self.config.joint_names) - points = [ - TrajectoryPoint( - positions=[float(value) for value in state], - velocities=zeros, - time_from_start=0.0, - ) - ] - points.extend( - TrajectoryPoint( - positions=[float(value) for value in action], - velocities=zeros, - time_from_start=(index + 1) / self.config.fps, - ) - for index, action in enumerate(actions) - ) - return JointTrajectory(joint_names=list(self.config.joint_names), points=points) + action_chunk = loaded.postprocessor(predict(prepared)) + result = np.asarray(action_chunk.to("cpu").numpy(), dtype=np.float32) + if result.ndim != 3 or result.shape[0] != 1: + raise RuntimeError(f"LeRobot returned invalid batched action shape {result.shape}") + return np.asarray(result[0], dtype=np.float32) + + def _require_loaded(self) -> _LoadedPolicy: + if self._loaded is None: + raise RuntimeError("LeRobot backend is not loaded") + return self._loaded + + +LeRobotPolicyRuntime = declare_policy_runtime( + "LeRobotPolicyRuntime", + __name__, + OpenYamLeRobotPolicy, + LeRobotBackend, +) - def _wait_for_newer_joint_state(self, previous_ts: float) -> None: - with self._observation_changed: - self._observation_changed.wait_for( - lambda: self._stop_event.is_set() - or ( - self._latest_joint_state is not None - and self._latest_joint_state.ts > previous_ts - ) - ) - def _cancel_trajectory(self) -> str | None: - try: - result = self._control.cancel_trajectory(task_name=self.config.trajectory_task_name) - except Exception as exc: - logger.exception( - "Failed to cancel policy trajectory", - task_name=self.config.trajectory_task_name, - ) - return f"Failed to cancel policy trajectory: {exc}" - if result.safe: - return None - message = result.message or "Policy trajectory cancellation was uncertain" - logger.error( - "Policy trajectory cancellation was uncertain", - error=message, - task_name=self.config.trajectory_task_name, +def _validate_features(policy_config: PreTrainedConfig, profile: PolicyIOProfile) -> None: + inputs = policy_config.input_features or {} + outputs = policy_config.output_features or {} + missing = set(profile.observations) - set(inputs) + if missing: + raise ValueError(f"LeRobot checkpoint is missing input features: {sorted(missing)}") + if profile.action.key not in outputs: + raise ValueError(f"LeRobot checkpoint has no {profile.action.key!r} output feature") + if getattr(policy_config, "temporal_ensemble_coeff", None) is not None: + raise ValueError("Policies using temporal ensembling are not supported") + + for key, source in profile.observations.items(): + actual = tuple(inputs[key].shape) + expected = ( + (source.shape[2], source.shape[0], source.shape[1]) + if isinstance(source, ImageSource) + else (len(source.joints),) ) - return message + if actual != expected: + raise ValueError(f"LeRobot feature {key!r} shape {actual} does not match {expected}") + action_shape = tuple(outputs[profile.action.key].shape) + expected_action = (len(profile.action.demonstration.joints),) + if action_shape != expected_action: + raise ValueError(f"LeRobot action shape {action_shape} does not match {expected_action}") def _checkpoint_action_bounds( @@ -546,19 +188,11 @@ def _checkpoint_action_bounds( break if lower_tensor is None or upper_tensor is None: raise ValueError("Policy postprocessor has no action min/max statistics") - lower = np.asarray(lower_tensor.detach().cpu().numpy(), dtype=np.float32) upper = np.asarray(upper_tensor.detach().cpu().numpy(), dtype=np.float32) - expected_shape = (expected_width,) - if lower.shape != expected_shape or upper.shape != expected_shape: - raise ValueError( - "Policy action range shape does not match configured joints: " - f"min={lower.shape}, max={upper.shape}, expected={expected_shape}" - ) - if not np.all(np.isfinite(lower)) or not np.all(np.isfinite(upper)): - raise ValueError("Policy action range contains non-finite values") - if np.any(lower > upper): - raise ValueError("Policy action range has min greater than max") + shape = (expected_width,) + if lower.shape != shape or upper.shape != shape: + raise ValueError(f"Policy action range shape must be {shape}") return lower, upper diff --git a/dimos/imitation/policy/lerobot/python/dimos_lerobot/runtime_tests.py b/dimos/imitation/policy/lerobot/python/dimos_lerobot/runtime_tests.py index c427b41fd6..406aebeb72 100644 --- a/dimos/imitation/policy/lerobot/python/dimos_lerobot/runtime_tests.py +++ b/dimos/imitation/policy/lerobot/python/dimos_lerobot/runtime_tests.py @@ -12,36 +12,19 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Behavior tests for the isolated LeRobot runtime.""" +from typing import Any -from collections.abc import Callable, Iterator -from threading import Event, Thread -import time -from typing import Any, Protocol - -from dimos_lerobot import runtime as policy_runtime -from dimos_lerobot.runtime import LeRobotPolicyRuntime -from lerobot.configs.policies import PreTrainedConfig -import numpy as np -from numpy.typing import NDArray +from dimos_lerobot.runtime import ( + LeRobotPolicyRuntime, + _checkpoint_action_bounds, + _validate_features, +) import pytest -import pytest_mock import torch -from torch import Tensor - -from dimos.control.tasks.trajectory_task.trajectory_task import ( - TrajectoryCancellationResult, - TrajectoryCancellationStatus, - TrajectoryExecutionResult, - TrajectoryExecutionStatus, -) -from dimos.msgs.sensor_msgs.Image import Image, ImageFormat -from dimos.msgs.sensor_msgs.JointState import JointState -from dimos.protocol.rpc.pubsubrpc import LCMRPC -from dimos.teleop.quest.quest_types import Buttons -from dimos.utils.testing.waiting import wait_until -JOINTS = [f"test_arm/joint{i}" for i in range(1, 5)] +from dimos.experimental.isolated_python.bootstrap import validate_runtime +from dimos.imitation.policy.lerobot.module import OpenYamLeRobotPolicy +from dimos.robot.manipulators.openyam.learning import OPENYAM_QUEST_IO class FakeFeature: @@ -49,611 +32,45 @@ def __init__(self, shape: tuple[int, ...]) -> None: self.shape = shape -class FakeUpstreamConfig: - def __init__(self, joint_count: int, n_action_steps: int) -> None: - self.type = "fake_policy" - self.device: str | None = "cpu" - self.use_amp = False - self.chunk_size = 3 - self.n_action_steps: int | None = n_action_steps - self.temporal_ensemble_coeff: float | None = None - self.input_features = { - "observation.images.wrist": FakeFeature((3, 4, 5)), - "observation.state": FakeFeature((joint_count,)), - } - self.output_features = {"action": FakeFeature((joint_count,))} - - -class FakePipeline: - def __init__(self, steps: list[object] | None = None) -> None: - self.calls: list[object] = [] - self.reset_count = 0 - self.steps = steps or [] - - def __call__(self, value: object) -> object: - self.calls.append(value) - return value +class FakeConfig: + temporal_ensemble_coeff = None + input_features = { + "observation.images.wrist": FakeFeature((3, 480, 640)), + "observation.state": FakeFeature((7,)), + } + output_features = {"action": FakeFeature((7,))} - def reset(self) -> None: - self.reset_count += 1 - -class FakeActionStats: - def __init__(self, lower: list[float], upper: list[float]) -> None: - self._state = { - "action.min": torch.tensor(lower), - "action.max": torch.tensor(upper), +class FakeStats: + def state_dict(self) -> dict[str, torch.Tensor]: + return { + "action.min": torch.zeros(7), + "action.max": torch.ones(7), } - def state_dict(self) -> dict[str, Tensor]: - return self._state - - -class FakePolicy: - def __init__( - self, - action_chunk: NDArray[np.float32], - n_action_steps: int = 2, - *, - action_lower: list[float] | None = None, - action_upper: list[float] | None = None, - ) -> None: - self.action_chunk = torch.from_numpy(action_chunk).unsqueeze(0) - self.called = Event() - self.reset_count = 0 - self.batch: dict[str, object] | None = None - self.upstream_config = FakeUpstreamConfig(len(JOINTS), n_action_steps) - self.preprocessor = FakePipeline() - self.postprocessor = FakePipeline( - [ - FakeActionStats( - action_lower or [-100.0] * len(JOINTS), - action_upper or [100.0] * len(JOINTS), - ) - ] - ) - self.config_load_count = 0 - - def reset(self) -> None: - self.reset_count += 1 - - def predict_action_chunk(self, batch: dict[str, object]) -> Tensor: - self.batch = dict(batch) - self.called.set() - return self.action_chunk - - -class RuntimeFactory(Protocol): - def __call__( - self, - policy: FakePolicy, - *, - device: str | None = None, - ) -> tuple[LeRobotPolicyRuntime, Any]: ... - - -@pytest.fixture -def make_runtime(mocker: pytest_mock.MockerFixture) -> Iterator[RuntimeFactory]: - mocker.patch("dimos.core.module.get_loop", return_value=(mocker.MagicMock(), None)) - mocker.patch.object(LCMRPC, "__init__", return_value=None) - mocker.patch.object(LCMRPC, "serve_module_rpc", return_value=None) - mocker.patch.object(LCMRPC, "start", return_value=None) - mocker.patch.object(LCMRPC, "stop", return_value=None) - built: list[LeRobotPolicyRuntime] = [] - - def _make( - policy: FakePolicy, - *, - device: str | None = None, - ) -> tuple[LeRobotPolicyRuntime, Any]: - def load_config(_path: str) -> FakeUpstreamConfig: - policy.config_load_count += 1 - return policy.upstream_config - - policy_class = mocker.MagicMock() - policy_class.from_pretrained.return_value = policy - mocker.patch.object(PreTrainedConfig, "from_pretrained", side_effect=load_config) - mocker.patch.object(policy_runtime, "get_policy_class", return_value=policy_class) - mocker.patch.object( - policy_runtime, - "make_pre_post_processors", - return_value=(policy.preprocessor, policy.postprocessor), - ) - - def prepare_observation( - observation: dict[str, NDArray[np.uint8] | NDArray[np.float32]], - _device: torch.device, - *, - task: str, - robot_type: str, - ) -> dict[str, object]: - return { - **{ - name: torch.from_numpy(value).unsqueeze(0) - for name, value in observation.items() - }, - "task": task, - "robot_type": robot_type, - } - - mocker.patch.object( - policy_runtime, - "prepare_observation_for_inference", - side_effect=prepare_observation, - ) - mocker.patch.object(policy_runtime, "register_third_party_plugins") - - module = LeRobotPolicyRuntime( - _isolated_python_runtime=True, - policy_path="checkpoint/default", - task="pick up the test object", - device=device, - joint_names=JOINTS, - fps=50.0, - robot_type="test_arm", - image_width=5, - image_height=4, - ) - control = mocker.MagicMock() - control.execute_trajectory.return_value = TrajectoryExecutionResult( - TrajectoryExecutionStatus.ACCEPTED - ) - control.cancel_trajectory.return_value = TrajectoryCancellationResult( - TrajectoryCancellationStatus.ALREADY_STOPPED - ) - control.list_tasks.return_value = ["policy_rollout"] - mocker.patch.object(module, "_control", control, create=True) - built.append(module) - return module, control - - yield _make - for module in built: - module.stop() - - -def _action_chunk(steps: int = 3) -> NDArray[np.float32]: - return np.arange(steps * len(JOINTS), dtype=np.float32).reshape( - steps, len(JOINTS) - ) / np.float32(10) - - -def _provide_observation( - module: LeRobotPolicyRuntime, - *, - positions: list[float] | None = None, - ts: float | None = None, -) -> tuple[NDArray[np.uint8], list[float], float]: - rgb = np.zeros((4, 5, 3), dtype=np.uint8) - rgb[..., 0] = 10 - rgb[..., 1] = 20 - rgb[..., 2] = 30 - values = positions or [float(i) / 10 for i in range(len(JOINTS))] - timestamp = time.time() if ts is None else ts - module._on_color_image(Image(data=rgb, format=ImageFormat.RGB, ts=timestamp)) - module._on_joint_state(JointState(ts=timestamp, name=JOINTS, position=values)) - return rgb, values, timestamp - - -def _preflight(module: LeRobotPolicyRuntime) -> None: - status = module.preflight_rollout() - assert status["policy_ready"] is True - assert status["observations_ready"] is True - assert status["last_error"] is None - - -def test_policy_predicts_and_executes_one_native_joint_chunk(make_runtime: RuntimeFactory) -> None: - actions = _action_chunk() - policy = FakePolicy(actions, n_action_steps=2) - module, control = make_runtime(policy) - rgb, positions, _ts = _provide_observation(module) - - _preflight(module) - assert module.start_rollout()["active"] is True - wait_until(lambda: control.execute_trajectory.call_count >= 1, timeout=1.0) - module.stop_rollout() - - call = control.execute_trajectory.call_args_list[0] - trajectory = call.args[0] - assert call.kwargs == {"task_name": "policy_rollout"} - assert trajectory.joint_names == JOINTS - assert [point.time_from_start for point in trajectory.points] == [0.0, 0.02, 0.04] - np.testing.assert_allclose(trajectory.points[0].positions, positions) - np.testing.assert_allclose(trajectory.points[1].positions, actions[0]) - np.testing.assert_allclose(trajectory.points[2].positions, actions[1]) - assert policy.batch is not None - assert policy.batch["task"] == "pick up the test object" - image = policy.batch["observation.images.wrist"] - assert isinstance(image, Tensor) - np.testing.assert_array_equal(image.squeeze(0).numpy(), rgb) - assert policy.postprocessor.calls - assert module.rollout_status()["chunks_accepted"] >= 1 - - -def test_policy_actions_are_clipped_to_checkpoint_range(make_runtime: RuntimeFactory) -> None: - actions = np.zeros((3, len(JOINTS)), dtype=np.float32) - actions[:, -1] = 1.016 - policy = FakePolicy( - actions, - n_action_steps=2, - action_lower=[-10.0, -10.0, -10.0, 0.0], - action_upper=[10.0, 10.0, 10.0, 1.0], - ) - module, control = make_runtime(policy) - _provide_observation(module) - - _preflight(module) - module.start_rollout() - wait_until(lambda: control.execute_trajectory.call_count >= 1, timeout=1.0) - module.stop_rollout() - - trajectory = control.execute_trajectory.call_args_list[0].args[0] - assert [point.positions[-1] for point in trajectory.points[1:]] == [1.0, 1.0] - - -def test_next_chunk_uses_latest_joint_observation(make_runtime: RuntimeFactory) -> None: - policy = FakePolicy(_action_chunk(), n_action_steps=1) - module, control = make_runtime(policy) - first_submitted = Event() - release_first = Event() - - def execute_trajectory(*_args: object, **_kwargs: object) -> TrajectoryExecutionResult: - if not first_submitted.is_set(): - first_submitted.set() - assert release_first.wait(timeout=1.0) - return TrajectoryExecutionResult(TrajectoryExecutionStatus.ACCEPTED) - - control.execute_trajectory.side_effect = execute_trajectory - _provide_observation(module, positions=[0.0] * len(JOINTS)) - _preflight(module) - module.start_rollout() - assert first_submitted.wait(timeout=1.0) - - latest = [0.4, 0.3, 0.2, 0.1] - _provide_observation(module, positions=latest) - release_first.set() - wait_until(lambda: control.execute_trajectory.call_count >= 2, timeout=1.0) - module.stop_rollout() - - second = control.execute_trajectory.call_args_list[1].args[0] - np.testing.assert_allclose(second.points[0].positions, latest) - - -def test_start_mismatch_waits_for_new_joint_state_before_retry( - make_runtime: RuntimeFactory, -) -> None: - policy = FakePolicy(_action_chunk(), n_action_steps=1) - module, control = make_runtime(policy) - _provide_observation(module) - - def execute_trajectory(*_args: object, **_kwargs: object) -> TrajectoryExecutionResult: - status = ( - TrajectoryExecutionStatus.START_STATE_MISMATCH - if control.execute_trajectory.call_count == 1 - else TrajectoryExecutionStatus.ACCEPTED - ) - return TrajectoryExecutionResult(status) - - control.execute_trajectory.side_effect = execute_trajectory - _preflight(module) - module.start_rollout() - wait_until(lambda: control.execute_trajectory.call_count == 1, timeout=1.0) - - _provide_observation(module, positions=[0.2] * len(JOINTS), ts=time.time() + 0.01) - wait_until(lambda: control.execute_trajectory.call_count >= 2, timeout=1.0) - module.stop_rollout() - - assert module.rollout_status()["last_error"] is None - - -def test_a_press_stops_worker_before_cancelling_its_trajectory( - make_runtime: RuntimeFactory, -) -> None: - policy = FakePolicy(_action_chunk(), n_action_steps=2) - module, control = make_runtime(policy) - _provide_observation(module) - _preflight(module) - execute_started = Event() - release_execute = Event() - stop_finished = Event() - - def execute_trajectory(*_args: object, **_kwargs: object) -> TrajectoryExecutionResult: - execute_started.set() - assert release_execute.wait(timeout=1.0) - return TrajectoryExecutionResult(TrajectoryExecutionStatus.ACCEPTED) - - control.execute_trajectory.side_effect = execute_trajectory - pressed = Buttons() - pressed.right_primary = True - def stop_from_button() -> None: - module._on_button_pressed(pressed) - stop_finished.set() - - module._on_button_pressed(pressed) - assert execute_started.wait(timeout=1.0) - stop_thread = Thread(target=stop_from_button) - stop_thread.start() - try: - assert module._stop_event.wait(timeout=1.0) - assert module.rollout_status()["active"] is True - control.cancel_trajectory.assert_not_called() - finally: - release_execute.set() - stop_thread.join(timeout=1.0) - - assert stop_finished.is_set() - assert control.execute_trajectory.call_count == 1 - control.cancel_trajectory.assert_called_with(task_name="policy_rollout") - - -def test_uncertain_cancellation_is_reported(make_runtime: RuntimeFactory) -> None: - policy = FakePolicy(_action_chunk(), n_action_steps=1) - module, control = make_runtime(policy) - _provide_observation(module) - _preflight(module) - control.cancel_trajectory.return_value = TrajectoryCancellationResult( - TrajectoryCancellationStatus.UNCERTAIN, - "coordinator did not confirm cancellation", - ) - - module.start_rollout() - wait_until(lambda: control.execute_trajectory.call_count >= 1, timeout=1.0) - status = module.stop_rollout() - - assert status["active"] is False - assert status["last_error"] == "coordinator did not confirm cancellation" - - -@pytest.mark.parametrize( - ("actions", "message"), - [ - (np.zeros((3, len(JOINTS) - 1), dtype=np.float32), "action width"), - (np.full((3, len(JOINTS)), np.nan, dtype=np.float32), "non-finite joint targets"), - ], -) -def test_invalid_action_chunk_cancels_and_latches_rollout_off( - make_runtime: RuntimeFactory, - actions: NDArray[np.float32], - message: str, -) -> None: - policy = FakePolicy(actions) - module, control = make_runtime(policy) - _provide_observation(module) - - _preflight(module) - module.start_rollout() - - wait_until(lambda: module.rollout_status()["active"] is False, timeout=1.0) - assert message in (module.rollout_status()["last_error"] or "") - control.cancel_trajectory.assert_called_with(task_name="policy_rollout") - - -def test_trajectory_rejection_cancels_and_latches_rollout_off( - make_runtime: RuntimeFactory, -) -> None: - policy = FakePolicy(_action_chunk()) - module, control = make_runtime(policy) - _provide_observation(module) - control.execute_trajectory.return_value = TrajectoryExecutionResult( - TrajectoryExecutionStatus.POSITION_LIMIT_VIOLATION, - "outside hardware limits", - ) - - _preflight(module) - module.start_rollout() - - wait_until(lambda: module.rollout_status()["active"] is False, timeout=1.0) - assert module.rollout_status()["last_error"] == "outside hardware limits" - control.cancel_trajectory.assert_called_with(task_name="policy_rollout") - - -@pytest.mark.parametrize( - ("configure", "message"), - [ - (lambda config: setattr(config, "n_action_steps", None), "positive int"), - (lambda config: setattr(config, "n_action_steps", 0), "positive int"), - ( - lambda config: setattr(config, "temporal_ensemble_coeff", 0.01), - "temporal ensembling", - ), - ( - lambda config: setattr( - config.input_features["observation.images.wrist"], - "shape", - (3, 8, 8), - ), - "Policy image shape", - ), - ], -) -def test_incompatible_chunk_contract_is_rejected( - make_runtime: RuntimeFactory, - configure: Callable[[FakeUpstreamConfig], None], - message: str, -) -> None: - policy = FakePolicy(_action_chunk()) - configure(policy.upstream_config) - module, control = make_runtime(policy) - _provide_observation(module) - - status = module.preflight_rollout() - - assert status["active"] is False - assert status["policy_ready"] is False - assert message in (status["last_error"] or "") - control.execute_trajectory.assert_not_called() - - -def test_policy_refuses_to_load_without_live_observations(make_runtime: RuntimeFactory) -> None: - policy = FakePolicy(_action_chunk()) - module, control = make_runtime(policy) - - result = module.preflight_rollout() - - assert result["active"] is False - assert "no camera image" in (result["last_error"] or "") - assert policy.config_load_count == 0 - control.execute_trajectory.assert_not_called() - - -def test_policy_refuses_stale_observations(make_runtime: RuntimeFactory) -> None: - policy = FakePolicy(_action_chunk()) - module, control = make_runtime(policy) - _provide_observation(module, ts=time.time() - module.config.max_observation_age_s - 1.0) - - result = module.preflight_rollout() - - assert result["active"] is False - assert "camera image is stale" in (result["last_error"] or "") - assert policy.config_load_count == 0 - control.execute_trajectory.assert_not_called() - - -def test_checkpoint_loads_on_demand_and_is_cached(make_runtime: RuntimeFactory) -> None: - policy = FakePolicy(_action_chunk()) - module, control = make_runtime(policy) - _provide_observation(module) - - _preflight(module) - module.start_rollout() - wait_until(lambda: control.execute_trajectory.call_count >= 1, timeout=1.0) - module.stop_rollout() - control.execute_trajectory.reset_mock() - module.start_rollout() - wait_until(lambda: control.execute_trajectory.call_count >= 1, timeout=1.0) - module.stop_rollout() - - assert policy.config_load_count == 1 - - -def test_start_requires_a_successful_preflight(make_runtime: RuntimeFactory) -> None: - policy = FakePolicy(_action_chunk()) - module, control = make_runtime(policy) - _provide_observation(module) - - status = module.start_rollout() - - assert status["active"] is False - assert status["last_error"] == "policy preflight has not passed" - control.execute_trajectory.assert_not_called() - - -def test_preflight_never_sends_a_trajectory(make_runtime: RuntimeFactory) -> None: - policy = FakePolicy(_action_chunk()) - module, control = make_runtime(policy) - _provide_observation(module) - - status = module.preflight_rollout() - - assert status["policy_ready"] is True - control.execute_trajectory.assert_not_called() - control.cancel_trajectory.assert_not_called() - - -def test_preflight_requires_the_configured_coordinator_task( - make_runtime: RuntimeFactory, -) -> None: - policy = FakePolicy(_action_chunk()) - module, control = make_runtime(policy) - _provide_observation(module) - control.list_tasks.return_value = ["another_task"] - - status = module.preflight_rollout() - - assert status["policy_ready"] is False - assert "missing configured rollout task" in (status["last_error"] or "") - assert policy.config_load_count == 0 - - -@pytest.mark.parametrize( - ("positions", "names", "message"), - [ - ([0.0, 0.0, 0.0, float("nan")], JOINTS, "non-finite positions"), - ([0.0, 0.0, 0.0], JOINTS[:-1], "missing configured joints"), - ], -) -def test_preflight_rejects_invalid_live_joints( - make_runtime: RuntimeFactory, - positions: list[float], - names: list[str], - message: str, -) -> None: - policy = FakePolicy(_action_chunk()) - module, control = make_runtime(policy) - timestamp = time.time() - module._on_color_image( - Image(data=np.zeros((4, 5, 3), dtype=np.uint8), format=ImageFormat.RGB, ts=timestamp) - ) - module._on_joint_state(JointState(ts=timestamp, name=names, position=positions)) - - status = module.preflight_rollout() - - assert status["policy_ready"] is False - assert message in (status["last_error"] or "") - control.execute_trajectory.assert_not_called() - - -@pytest.mark.parametrize( - ("image", "image_format"), - [ - (np.zeros((8, 8, 3), dtype=np.uint8), ImageFormat.RGB), - (np.zeros((4, 5, 3), dtype=np.uint8), ImageFormat.BGR), - ], -) -def test_preflight_requires_exact_live_rgb_contract( - make_runtime: RuntimeFactory, - image: NDArray[np.uint8], - image_format: ImageFormat, -) -> None: - policy = FakePolicy(_action_chunk()) - module, control = make_runtime(policy) - timestamp = time.time() - module._on_color_image(Image(data=image, format=image_format, ts=timestamp)) - module._on_joint_state(JointState(ts=timestamp, name=JOINTS, position=[0.0] * len(JOINTS))) +class FakePipeline: + steps: list[Any] = [FakeStats()] - status = module.preflight_rollout() - assert status["policy_ready"] is False - assert "no camera image" in (status["last_error"] or "") - control.execute_trajectory.assert_not_called() +def test_generated_runtime_implements_the_host_contract() -> None: + validate_runtime(OpenYamLeRobotPolicy, LeRobotPolicyRuntime) -@pytest.mark.parametrize( - ("lower", "upper", "message"), - [ - ([-1.0] * 3, [1.0] * 3, "range shape"), - ([-1.0, -1.0, -1.0, float("nan")], [1.0] * 4, "non-finite"), - ([2.0] * 4, [1.0] * 4, "min greater than max"), - ], -) -def test_preflight_rejects_invalid_checkpoint_action_bounds( - make_runtime: RuntimeFactory, - lower: list[float], - upper: list[float], - message: str, -) -> None: - policy = FakePolicy(_action_chunk(), action_lower=lower, action_upper=upper) - module, control = make_runtime(policy) - _provide_observation(module) +def test_lerobot_feature_validation_uses_profile_keys_and_shapes() -> None: + _validate_features(FakeConfig(), OPENYAM_QUEST_IO) # type: ignore[arg-type] - status = module.preflight_rollout() - assert status["policy_ready"] is False - assert message in (status["last_error"] or "") - control.execute_trajectory.assert_not_called() +def test_lerobot_feature_validation_rejects_missing_profile_key() -> None: + config = FakeConfig() + config.input_features = {"observation.state": FakeFeature((7,))} + with pytest.raises(ValueError, match="observation.images.wrist"): + _validate_features(config, OPENYAM_QUEST_IO) # type: ignore[arg-type] -def test_preflight_rejects_unavailable_cuda( - make_runtime: RuntimeFactory, - mocker: pytest_mock.MockerFixture, -) -> None: - policy = FakePolicy(_action_chunk()) - module, control = make_runtime(policy, device="cuda") - _provide_observation(module) - mocker.patch.object(torch.cuda, "is_available", return_value=False) - status = module.preflight_rollout() +def test_lerobot_action_bounds_are_extracted_for_common_safety_loop() -> None: + lower, upper = _checkpoint_action_bounds(FakePipeline(), 7) - assert status["policy_ready"] is False - assert "CUDA is not available" in (status["last_error"] or "") - control.execute_trajectory.assert_not_called() + assert lower.tolist() == [0.0] * 7 + assert upper.tolist() == [1.0] * 7 diff --git a/dimos/imitation/policy/lerobot/python/pyproject.toml b/dimos/imitation/policy/lerobot/python/pyproject.toml index 83c9c7a055..ef489674dd 100644 --- a/dimos/imitation/policy/lerobot/python/pyproject.toml +++ b/dimos/imitation/policy/lerobot/python/pyproject.toml @@ -46,6 +46,10 @@ strict = true explicit_package_bases = true mypy_path = "../../../../../" +[[tool.mypy.overrides]] +module = ["dimos", "dimos.*"] +follow_imports = "skip" + [[tool.mypy.overrides]] module = [ "lerobot.configs.policies", diff --git a/dimos/imitation/policy/lerobot/test_module.py b/dimos/imitation/policy/lerobot/test_module.py index 1395336f70..8d4cbcb698 100644 --- a/dimos/imitation/policy/lerobot/test_module.py +++ b/dimos/imitation/policy/lerobot/test_module.py @@ -18,15 +18,16 @@ import pytest from dimos.experimental.isolated_python.module import contract_rpc_names -from dimos.imitation.policy.lerobot.module import ( - LeRobotPolicyModule, - LeRobotPolicyModuleConfig, -) +from dimos.imitation.policy.lerobot.module import LeRobotPolicyConfig, OpenYamLeRobotPolicy -def test_contract_imports_without_runtime_dependencies() -> None: - assert LeRobotPolicyModule.implementation == "dimos_lerobot.runtime:LeRobotPolicyRuntime" - assert contract_rpc_names(LeRobotPolicyModule) == { +def test_generated_contract_has_profile_ports_and_rpc_surface() -> None: + blueprint = OpenYamLeRobotPolicy.blueprint(artifact="unused", task="test") + streams = {stream.name for stream in blueprint.blueprints[0].streams} + + assert OpenYamLeRobotPolicy.implementation == ("dimos_lerobot.runtime:LeRobotPolicyRuntime") + assert streams == {"button_pressed", "wrist_image", "coordinator_joint_state"} + assert contract_rpc_names(OpenYamLeRobotPolicy) == { "preflight_rollout", "rollout_status", "start_rollout", @@ -35,63 +36,27 @@ def test_contract_imports_without_runtime_dependencies() -> None: def test_contract_resolves_sibling_runtime_project() -> None: - module = LeRobotPolicyModule( - policy_path="unused", - task="test task", - joint_names=["joint"], - ) + module = OpenYamLeRobotPolicy(artifact="unused", task="test task") try: assert module.runtime_project == Path(__file__).parent / "python" finally: module.stop() -@pytest.mark.parametrize( - ("config", "message"), - [ - ( - { - "policy_path": "checkpoint", - "task": "test task", - "joint_names": ["joint1", "joint1"], - }, - "joint_names must not contain duplicates", - ), - ( - { - "policy_path": " ", - "task": "test task", - "joint_names": ["joint1"], - }, - "policy_path must not be blank", - ), - ( - { - "policy_path": "checkpoint", - "task": "test task", - "joint_names": ["joint1"], - "rollout_button": "NOPE", - }, - "unknown Quest button", - ), - ], -) -def test_config_rejects_ambiguous_names(config: dict[str, object], message: str) -> None: - with pytest.raises(ValidationError, match=message): - LeRobotPolicyModuleConfig.model_validate(config) +def test_config_rejects_blank_artifact_and_unknown_button() -> None: + with pytest.raises(ValidationError, match="artifact must not be blank"): + LeRobotPolicyConfig(artifact=" ", task="test") + with pytest.raises(ValidationError, match="unknown Quest button"): + LeRobotPolicyConfig(artifact="checkpoint", task="test", rollout_button="NOPE") -def test_existing_relative_checkpoint_is_resolved_before_isolation( +def test_existing_relative_artifact_is_resolved( tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: checkpoint = tmp_path / "checkpoint" checkpoint.mkdir() monkeypatch.chdir(tmp_path) - config = LeRobotPolicyModuleConfig( - policy_path="checkpoint", - task="test task", - joint_names=["joint1"], - ) + config = LeRobotPolicyConfig(artifact="checkpoint", task="test task") - assert config.policy_path == str(checkpoint) + assert config.artifact == str(checkpoint) diff --git a/dimos/imitation/policy/module.py b/dimos/imitation/policy/module.py new file mode 100644 index 0000000000..0bf67486c7 --- /dev/null +++ b/dimos/imitation/policy/module.py @@ -0,0 +1,161 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Profile-driven host contract for isolated policy rollout.""" + +from __future__ import annotations + +from pathlib import Path +from typing import ClassVar, Protocol, TypedDict + +from pydantic import Field, field_validator + +from dimos.control.tasks.trajectory_task.trajectory_task import ( + TrajectoryCancellationResult, + TrajectoryExecutionResult, +) +from dimos.core.core import rpc +from dimos.core.stream import In +from dimos.experimental.isolated_python.module import ( + IsolatedPythonModule, + IsolatedPythonModuleConfig, +) +from dimos.imitation.profile import ImageSource, PolicyIOProfile +from dimos.msgs.sensor_msgs.Image import Image +from dimos.msgs.sensor_msgs.JointState import JointState +from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory +from dimos.spec.utils import Spec +from dimos.teleop.quest.quest_types import BUTTON_ALIASES, Buttons + +POLICY_ROLLOUT_TASK_NAME = "policy_rollout" +POLICY_ROLLOUT_INSTANCE_NAME = "PolicyRolloutModule" + + +class PolicyControlSpec(Spec, Protocol): + """Coordinator operations used by policy rollout.""" + + def execute_trajectory( + self, + trajectory: JointTrajectory, + task_name: str, + ) -> TrajectoryExecutionResult: ... + + def cancel_trajectory(self, task_name: str) -> TrajectoryCancellationResult: ... + + def list_tasks(self) -> list[str]: ... + + +class RolloutStatus(TypedDict): + """Operator-facing state of the configured policy rollout.""" + + active: bool + artifact: str + backend: str | None + task: str + device: str | None + policy_ready: bool + observations_ready: bool + chunks_accepted: int + last_error: str | None + + +class PolicyRolloutConfig(IsolatedPythonModuleConfig): + """Backend-neutral rollout and safety configuration.""" + + artifact: str = Field(min_length=1) + task: str = Field(min_length=1) + device: str | None = None + max_observation_age_s: float = Field(default=0.5, gt=0) + max_execution_horizon_s: float = Field(default=0.5, gt=0) + trajectory_task_name: str = POLICY_ROLLOUT_TASK_NAME + rollout_button: str = "A" + + @field_validator("artifact") + @classmethod + def artifact_must_not_be_blank(cls, artifact: str) -> str: + if not artifact.strip(): + raise ValueError("artifact must not be blank") + path = Path(artifact).expanduser() + return str(path.resolve()) if path.exists() else artifact + + @field_validator("trajectory_task_name") + @classmethod + def trajectory_task_name_must_not_be_blank(cls, name: str) -> str: + if not name.strip(): + raise ValueError("trajectory_task_name must not be blank") + return name + + @field_validator("rollout_button") + @classmethod + def rollout_button_must_be_digital(cls, name: str) -> str: + if BUTTON_ALIASES.get(name, name) not in Buttons.BITS: + raise ValueError(f"unknown Quest button {name!r}") + return name + + +class _PolicyModule(IsolatedPythonModule): + """RPC surface shared by all generated policy module declarations.""" + + config: PolicyRolloutConfig + profile: ClassVar[PolicyIOProfile] + button_pressed: In[Buttons] + _control: PolicyControlSpec + + @rpc + def preflight_rollout(self) -> RolloutStatus: + """Load and validate the policy and live inputs without moving the robot.""" + raise NotImplementedError + + @rpc + def start_rollout(self) -> RolloutStatus: + """Start the configured policy until explicitly stopped or it fails.""" + raise NotImplementedError + + @rpc + def stop_rollout(self) -> RolloutStatus: + """Stop rollout publication and clear the policy action queue.""" + raise NotImplementedError + + @rpc + def rollout_status(self) -> RolloutStatus: + """Return the lifecycle and observation state of the configured policy.""" + raise NotImplementedError + + +def declare_policy_module( + name: str, + module_name: str, + profile: PolicyIOProfile, + config_type: type[PolicyRolloutConfig], + implementation: str, +) -> type[_PolicyModule]: + """Declare a stable importable policy module with profile-shaped ports.""" + annotations: dict[str, object] = {"config": config_type} + for source in profile.observations.values(): + annotations[source.stream] = ( + In[Image] if isinstance(source, ImageSource) else In[JointState] + ) + + return type( + name, + (_PolicyModule,), + { + "__annotations__": annotations, + "__doc__": f"Policy rollout module for the {profile.name!r} profile.", + "__module__": module_name, + "__qualname__": name, + "implementation": implementation, + "profile": profile, + }, + ) diff --git a/dimos/imitation/policy/runtime.py b/dimos/imitation/policy/runtime.py new file mode 100644 index 0000000000..c3a4cc604e --- /dev/null +++ b/dimos/imitation/policy/runtime.py @@ -0,0 +1,485 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Backend-neutral live observation alignment and safe rollout execution.""" + +from __future__ import annotations + +from collections import deque +from collections.abc import Callable +from dataclasses import dataclass +import math +from threading import Condition, Event, RLock, Thread, current_thread +import time +from typing import Any, cast + +import numpy as np +from numpy.typing import NDArray +from reactivex.disposable import Disposable + +from dimos.constants import DEFAULT_THREAD_JOIN_TIMEOUT +from dimos.control.tasks.trajectory_task.trajectory_task import TrajectoryExecutionStatus +from dimos.core.core import rpc +from dimos.core.module import Module +from dimos.imitation.policy.backend import PolicyBackend, PolicyBackendInfo +from dimos.imitation.policy.module import ( + PolicyControlSpec, + PolicyRolloutConfig, + RolloutStatus, + _PolicyModule, +) +from dimos.imitation.profile import ImageSource, JointPositionSource, PolicyIOProfile +from dimos.msgs.sensor_msgs.Image import Image, ImageFormat +from dimos.msgs.sensor_msgs.JointState import JointState +from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory +from dimos.msgs.trajectory_msgs.TrajectoryPoint import TrajectoryPoint +from dimos.teleop.quest.quest_types import BUTTON_ALIASES, Buttons +from dimos.utils.logging_config import setup_logger + +logger = setup_logger() + + +@dataclass(frozen=True) +class _TimedValue: + value: NDArray[Any] + ts: float + + +class _PolicyRuntimeMixin: + """Runtime implementation mixed into a generated backend declaration.""" + + profile: PolicyIOProfile + backend_type: type[PolicyBackend] + config: PolicyRolloutConfig + button_pressed: Any + _control: PolicyControlSpec + register_disposable: Callable[[Any], None] + + def __init__(self, **kwargs: Any) -> None: + super().__init__(**kwargs) + self._lock = RLock() + self._observation_changed = Condition(self._lock) + self._buffers: dict[str, deque[_TimedValue]] = { + key: deque(maxlen=64) for key in self.profile.observations + } + self._backend = self.backend_type(self.config) + self._backend_info: PolicyBackendInfo | None = None + self._stop_event = Event() + self._thread: Thread | None = None + self._chunks_accepted = 0 + self._last_error: str | None = None + self._active = False + + @rpc + def start(self) -> None: + Module.start(cast("Module", self)) + streams = {source.stream for source in self.profile.observations.values()} + for stream_name in streams: + stream = getattr(self, stream_name) + self.register_disposable( + Disposable( + stream.subscribe( + lambda message, name=stream_name: self._on_observation(name, message) + ) + ) + ) + self.register_disposable(Disposable(self.button_pressed.subscribe(self._on_button_pressed))) + + @rpc + def stop(self) -> None: + if not self._stop_policy(): + self._cancel_after_stop_timeout() + Module.stop(cast("Module", self)) + + @rpc + def preflight_rollout(self) -> RolloutStatus: + with self._lock: + if self._active: + self._last_error = "cannot preflight while a policy rollout is active" + return self._status_locked() + try: + self._snapshot_observation(time.time()) + except Exception as exc: + self._backend_info = None + self._last_error = str(exc) + return self._status_locked() + + try: + tasks = set(self._control.list_tasks()) + if self.config.trajectory_task_name not in tasks: + raise RuntimeError( + "ControlCoordinator is missing configured rollout task " + f"{self.config.trajectory_task_name!r}" + ) + backend_info = self._backend_info or self._backend.load(self.profile) + self._validate_backend_info(backend_info) + with self._lock: + self._backend_info = backend_info + self._snapshot_observation(time.time()) + self._last_error = None + return self._status_locked() + except Exception as exc: + with self._lock: + self._backend_info = None + self._last_error = str(exc) + return self._status_locked() + + @rpc + def start_rollout(self) -> RolloutStatus: + with self._lock: + if self._thread is not None and self._thread.is_alive(): + self._last_error = "a policy rollout is already active" + return self._status_locked() + if self._backend_info is None: + self._last_error = "policy preflight has not passed" + return self._status_locked() + try: + self._snapshot_observation(time.time()) + except RuntimeError as exc: + self._last_error = str(exc) + return self._status_locked() + + self._stop_event.clear() + self._chunks_accepted = 0 + self._last_error = None + self._active = True + self._thread = Thread( + target=self._run_rollout, + name="policy-rollout", + daemon=True, + ) + self._thread.start() + return self._status_locked() + + @rpc + def stop_rollout(self) -> RolloutStatus: + if not self._stop_policy(): + self._cancel_after_stop_timeout() + return self.rollout_status() + + @rpc + def rollout_status(self) -> RolloutStatus: + with self._lock: + return self._status_locked() + + def _status_locked(self) -> RolloutStatus: + try: + self._snapshot_observation(time.time()) + observations_ready = True + except RuntimeError: + observations_ready = False + return { + "active": self._active, + "artifact": self.config.artifact, + "backend": self._backend_info.name if self._backend_info is not None else None, + "task": self.config.task, + "device": self.config.device, + "policy_ready": self._backend_info is not None, + "observations_ready": observations_ready, + "chunks_accepted": self._chunks_accepted, + "last_error": self._last_error, + } + + def _on_observation(self, stream_name: str, message: object) -> None: + changed_action_state = False + with self._observation_changed: + for key, source in self.profile.observations.items(): + if source.stream != stream_name: + continue + try: + value, ts = _read_source(source, message) + except (TypeError, ValueError) as exc: + logger.warning( + "Ignoring invalid policy observation", + feature=key, + error=str(exc), + ) + continue + self._buffers[key].append(_TimedValue(value=value, ts=ts)) + changed_action_state = changed_action_state or key == self.profile.action_state_key + if changed_action_state: + self._observation_changed.notify_all() + + def _on_button_pressed(self, buttons: Buttons) -> None: + button = BUTTON_ALIASES.get(self.config.rollout_button, self.config.rollout_button) + if not bool(getattr(buttons, button)): + return + with self._lock: + active = self._active + if active: + self.stop_rollout() + else: + self.start_rollout() + + def _snapshot_observation( + self, + now: float, + ) -> tuple[dict[str, NDArray[Any]], NDArray[np.float32], float]: + anchor_key = self.profile.sync.anchor + anchor_buffer = self._buffers[anchor_key] + if not anchor_buffer: + raise RuntimeError(f"no {anchor_key!r} observation has been received") + anchor = anchor_buffer[-1] + selected = {anchor_key: anchor} + tolerance_s = self.profile.sync.tolerance_ms / 1000.0 + + for key, buffer in self._buffers.items(): + if key == anchor_key: + continue + if not buffer: + raise RuntimeError(f"no {key!r} observation has been received") + nearest = min(buffer, key=lambda item: abs(item.ts - anchor.ts)) + skew = abs(nearest.ts - anchor.ts) + if skew > tolerance_s: + raise RuntimeError( + f"{key!r} is {skew * 1000.0:.1f}ms from anchor {anchor_key!r}; " + f"limit is {self.profile.sync.tolerance_ms:.1f}ms" + ) + selected[key] = nearest + + for key, item in selected.items(): + age = now - item.ts + if age > self.config.max_observation_age_s: + raise RuntimeError(f"{key!r} observation is stale by {age:.2f}s") + + state_item = selected[self.profile.action_state_key] + observations = {key: item.value.copy() for key, item in selected.items()} + return observations, np.asarray(state_item.value, dtype=np.float32), state_item.ts + + def _run_rollout(self) -> None: + info: PolicyBackendInfo | None = None + try: + with self._lock: + info = self._backend_info + if info is None: + raise RuntimeError("policy preflight has not passed") + if self._stop_event.is_set(): + return + self._backend.reset() + execution_steps = self._execution_steps(info) + + while not self._stop_event.is_set(): + with self._lock: + observations, state, state_ts = self._snapshot_observation(time.time()) + action_chunk = np.asarray( + self._backend.predict(observations, self.config.task), + dtype=np.float32, + ) + actions = self._validated_actions(action_chunk, info, execution_steps) + if self._stop_event.is_set(): + break + result = self._control.execute_trajectory( + self._trajectory(state, actions), + task_name=self.config.trajectory_task_name, + ) + if result.status is TrajectoryExecutionStatus.START_STATE_MISMATCH: + self._wait_for_newer_joint_state(state_ts) + continue + if result.status is not TrajectoryExecutionStatus.ACCEPTED: + raise RuntimeError( + result.message or f"trajectory rejected: {result.status.name}" + ) + with self._lock: + self._chunks_accepted += 1 + self._stop_event.wait(execution_steps / self.profile.sync.rate_hz) + except Exception as exc: + with self._lock: + self._last_error = str(exc) + logger.exception("Policy execution stopped", error=str(exc)) + finally: + self._stop_event.set() + cancellation_error = self._cancel_trajectory() + reset_error = self._reset_backend() if info is not None else None + with self._lock: + shutdown_errors = [ + error for error in (cancellation_error, reset_error) if error is not None + ] + if shutdown_errors: + shutdown_error = "; ".join(shutdown_errors) + self._last_error = ( + f"{self._last_error}; {shutdown_error}" + if self._last_error is not None + else shutdown_error + ) + self._active = False + + def _execution_steps(self, info: PolicyBackendInfo) -> int: + horizon_steps = math.floor( + self.config.max_execution_horizon_s * self.profile.sync.rate_hz + 1e-9 + ) + if horizon_steps < 1: + raise ValueError("max execution horizon is shorter than one policy step") + return min(info.chunk_length, info.preferred_execution_steps, horizon_steps) + + def _validated_actions( + self, + action_chunk: NDArray[np.float32], + info: PolicyBackendInfo, + execution_steps: int, + ) -> NDArray[np.float32]: + width = len(self.profile.action.demonstration.joints) + if action_chunk.ndim != 2 or action_chunk.shape[1] != width: + raise RuntimeError( + f"policy returned action chunk shape {action_chunk.shape}, expected (steps, {width})" + ) + if action_chunk.shape[0] < execution_steps: + raise RuntimeError( + f"policy returned {action_chunk.shape[0]} steps, expected at least {execution_steps}" + ) + actions = action_chunk[:execution_steps] + if not np.all(np.isfinite(actions)): + raise RuntimeError("policy returned non-finite joint targets") + if info.action_lower is None: + return actions + assert info.action_upper is not None + bounded = np.clip(actions, info.action_lower, info.action_upper) + if np.any(actions != bounded): + logger.warning("Clipped policy actions to backend range") + return bounded + + def _validate_backend_info(self, info: PolicyBackendInfo) -> None: + if info.chunk_length <= 0 or info.preferred_execution_steps <= 0: + raise ValueError("backend chunk and execution step counts must be positive") + bounds = (info.action_lower, info.action_upper) + if (bounds[0] is None) != (bounds[1] is None): + raise ValueError("backend must provide both action bounds or neither") + if bounds[0] is None: + return + assert bounds[1] is not None + shape = (len(self.profile.action.demonstration.joints),) + if bounds[0].shape != shape or bounds[1].shape != shape: + raise ValueError(f"backend action bounds must have shape {shape}") + if not np.all(np.isfinite(bounds[0])) or not np.all(np.isfinite(bounds[1])): + raise ValueError("backend action bounds contain non-finite values") + if np.any(bounds[0] > bounds[1]): + raise ValueError("backend action lower bound exceeds upper bound") + + def _stop_policy(self) -> bool: + with self._lock: + thread = self._thread + self._stop_event.set() + self._observation_changed.notify_all() + if thread is not None and thread is not current_thread(): + thread.join(timeout=DEFAULT_THREAD_JOIN_TIMEOUT) + return thread is None or not thread.is_alive() + + def _cancel_after_stop_timeout(self) -> None: + timeout_error = f"policy rollout did not stop within {DEFAULT_THREAD_JOIN_TIMEOUT} seconds" + cancellation_error = self._cancel_trajectory() + with self._lock: + self._last_error = ( + f"{timeout_error}; {cancellation_error}" + if cancellation_error is not None + else timeout_error + ) + + def _trajectory( + self, + state: NDArray[np.float32], + actions: NDArray[np.float32], + ) -> JointTrajectory: + joints = list(self.profile.action.demonstration.joints) + zeros = [0.0] * len(joints) + points = [ + TrajectoryPoint( + positions=[float(value) for value in state], + velocities=zeros, + time_from_start=0.0, + ) + ] + points.extend( + TrajectoryPoint( + positions=[float(value) for value in action], + velocities=zeros, + time_from_start=(index + 1) / self.profile.sync.rate_hz, + ) + for index, action in enumerate(actions) + ) + return JointTrajectory(joint_names=joints, points=points) + + def _wait_for_newer_joint_state(self, previous_ts: float) -> None: + state_buffer = self._buffers[self.profile.action_state_key] + with self._observation_changed: + self._observation_changed.wait_for( + lambda: self._stop_event.is_set() + or (bool(state_buffer) and state_buffer[-1].ts > previous_ts) + ) + + def _cancel_trajectory(self) -> str | None: + try: + result = self._control.cancel_trajectory(task_name=self.config.trajectory_task_name) + except Exception as exc: + logger.exception( + "Failed to cancel policy trajectory", + task_name=self.config.trajectory_task_name, + ) + return f"Failed to cancel policy trajectory: {exc}" + if result.safe: + return None + return result.message or "Policy trajectory cancellation was uncertain" + + def _reset_backend(self) -> str | None: + try: + self._backend.reset() + except Exception as exc: + logger.exception("Failed to reset policy backend") + return f"Failed to reset policy backend: {exc}" + return None + + +def declare_policy_runtime( + name: str, + module_name: str, + declaration: type[_PolicyModule], + backend_type: type[PolicyBackend], +) -> type[_PolicyModule]: + """Declare an importable runtime subclass using the common rollout loop.""" + return type( + name, + (_PolicyRuntimeMixin, declaration), + { + "__module__": module_name, + "__qualname__": name, + "backend_type": backend_type, + }, + ) + + +def _read_source( + source: ImageSource | JointPositionSource, + message: object, +) -> tuple[NDArray[Any], float]: + if isinstance(source, ImageSource): + if not isinstance(message, Image): + raise TypeError(f"expected Image, got {type(message).__name__}") + if message.format != ImageFormat.RGB or message.data.dtype != np.uint8: + raise ValueError("image must be uint8 RGB") + if message.data.shape != source.shape: + raise ValueError(f"image shape {message.data.shape} does not match {source.shape}") + return np.ascontiguousarray(message.data), message.ts + + if not isinstance(message, JointState): + raise TypeError(f"expected JointState, got {type(message).__name__}") + if len(message.name) != len(message.position): + raise ValueError("JointState names and positions have different lengths") + if len(message.name) != len(set(message.name)): + raise ValueError("JointState contains duplicate joint names") + positions = dict(zip(message.name, message.position, strict=True)) + missing = [joint for joint in source.joints if joint not in positions] + if missing: + raise ValueError(f"JointState is missing configured joints: {missing}") + value = np.asarray([positions[joint] for joint in source.joints], dtype=np.float32) + if not np.all(np.isfinite(value)): + raise ValueError("JointState contains non-finite positions") + return value, message.ts diff --git a/dimos/imitation/policy/test_runtime.py b/dimos/imitation/policy/test_runtime.py new file mode 100644 index 0000000000..1cb964d005 --- /dev/null +++ b/dimos/imitation/policy/test_runtime.py @@ -0,0 +1,320 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from collections.abc import Iterator, Mapping +from pathlib import Path +import time +from typing import Any + +import numpy as np +from numpy.typing import NDArray +import pytest +import pytest_mock + +from dimos.control.tasks.trajectory_task.trajectory_task import ( + TrajectoryCancellationResult, + TrajectoryCancellationStatus, + TrajectoryExecutionResult, + TrajectoryExecutionStatus, +) +from dimos.imitation.dataprep.core import SyncConfig +from dimos.imitation.policy.backend import PolicyBackendInfo +from dimos.imitation.policy.module import PolicyRolloutConfig, declare_policy_module +from dimos.imitation.policy.runtime import declare_policy_runtime +from dimos.imitation.profile import ( + ImageSource, + JointPositionAction, + JointPositionSource, + PolicyIOProfile, +) +from dimos.msgs.sensor_msgs.Image import Image, ImageFormat +from dimos.msgs.sensor_msgs.JointState import JointState +from dimos.protocol.rpc.pubsubrpc import LCMRPC +from dimos.utils.testing.waiting import wait_until + +JOINTS = ("left", "right") +PROFILE = PolicyIOProfile( + name="runtime-test", + robot_type="test", + observations={ + "top": ImageSource(stream="top_image", shape=(4, 5, 3)), + "left": ImageSource(stream="left_image", shape=(4, 5, 3)), + "state": JointPositionSource(stream="joint_state", joints=JOINTS), + }, + action=JointPositionAction( + key="actions", + demonstration=JointPositionSource(stream="joint_command", joints=JOINTS), + ), + sync=SyncConfig(anchor="top", rate_hz=30.0, tolerance_ms=20.0), +) + + +class _TestConfig(PolicyRolloutConfig): + pass + + +PolicyDeclaration = declare_policy_module( + "RuntimeTestPolicy", + __name__, + PROFILE, + _TestConfig, + "unused:TestRuntime", +) + + +class FakeBackend: + actions = np.arange(60, dtype=np.float32).reshape(30, 2) + + def __init__(self, _config: _TestConfig) -> None: + self.actions = type(self).actions.copy() + self.info = PolicyBackendInfo( + name="fake", + chunk_length=30, + preferred_execution_steps=20, + ) + self.load_count = 0 + self.reset_count = 0 + self.reset_error: str | None = None + + def load(self, _profile: PolicyIOProfile) -> PolicyBackendInfo: + self.load_count += 1 + return self.info + + def reset(self) -> None: + self.reset_count += 1 + if self.reset_error is not None: + raise RuntimeError(self.reset_error) + + def predict( + self, + _observations: Mapping[str, NDArray[Any]], + _task: str, + ) -> NDArray[np.float32]: + return self.actions.copy() + + +RuntimeModule = declare_policy_runtime( + "RuntimeTestModule", + __name__, + PolicyDeclaration, + FakeBackend, +) + + +@pytest.fixture +def runtime( + tmp_path: Path, + mocker: pytest_mock.MockerFixture, +) -> Iterator[tuple[RuntimeModule, Any]]: + mocker.patch("dimos.core.module.get_loop", return_value=(mocker.MagicMock(), None)) + mocker.patch.object(LCMRPC, "__init__", return_value=None) + mocker.patch.object(LCMRPC, "serve_module_rpc", return_value=None) + mocker.patch.object(LCMRPC, "start", return_value=None) + mocker.patch.object(LCMRPC, "stop", return_value=None) + module = RuntimeModule( + _isolated_python_runtime=True, + artifact=str(tmp_path / "artifact"), + task="test task", + ) + control = mocker.MagicMock() + control.list_tasks.return_value = ["policy_rollout"] + control.execute_trajectory.return_value = TrajectoryExecutionResult( + TrajectoryExecutionStatus.ACCEPTED + ) + control.cancel_trajectory.return_value = TrajectoryCancellationResult( + TrajectoryCancellationStatus.ALREADY_STOPPED + ) + mocker.patch.object(module, "_control", control, create=True) + yield module, control + module.stop() + + +def _image(ts: float) -> Image: + return Image(data=np.zeros((4, 5, 3), dtype=np.uint8), format=ImageFormat.RGB, ts=ts) + + +def _provide(module: RuntimeModule, *, top_ts: float, other_ts: float) -> None: + module._on_observation("top_image", _image(top_ts)) + module._on_observation("left_image", _image(other_ts)) + module._on_observation( + "joint_state", + JointState(ts=other_ts, name=list(reversed(JOINTS)), position=[2.0, 1.0]), + ) + + +def test_preflight_requires_every_camera_within_profile_tolerance( + runtime: tuple[RuntimeModule, Any], +) -> None: + module, _control = runtime + now = time.time() + _provide(module, top_ts=now, other_ts=now - 0.021) + + status = module.preflight_rollout() + + assert status["policy_ready"] is False + assert status["observations_ready"] is False + assert "21.0ms" in (status["last_error"] or "") + + +def test_rollout_caps_execution_horizon_and_preserves_profile_joint_order( + runtime: tuple[RuntimeModule, Any], +) -> None: + module, control = runtime + now = time.time() + _provide(module, top_ts=now, other_ts=now - 0.005) + assert module.preflight_rollout()["policy_ready"] is True + + module.start_rollout() + wait_until(lambda: control.execute_trajectory.call_count >= 1, timeout=1.0) + module.stop_rollout() + + trajectory = control.execute_trajectory.call_args.args[0] + assert trajectory.joint_names == list(JOINTS) + assert len(trajectory.points) == 16 + assert trajectory.points[-1].time_from_start == pytest.approx(0.5) + assert trajectory.points[0].positions == [1.0, 2.0] + np.testing.assert_array_equal(trajectory.points[-1].positions, FakeBackend.actions[14]) + + +def test_missing_declared_top_camera_fails_before_backend_load( + runtime: tuple[RuntimeModule, Any], +) -> None: + module, _control = runtime + now = time.time() + module._on_observation("left_image", _image(now)) + module._on_observation( + "joint_state", + JointState(ts=now, name=list(JOINTS), position=[1.0, 2.0]), + ) + + status = module.preflight_rollout() + + assert status["policy_ready"] is False + assert status["last_error"] == "no 'top' observation has been received" + assert module._backend.load_count == 0 + + +def test_start_requires_successful_preflight( + runtime: tuple[RuntimeModule, Any], +) -> None: + module, control = runtime + now = time.time() + _provide(module, top_ts=now, other_ts=now) + + status = module.start_rollout() + + assert status["active"] is False + assert status["last_error"] == "policy preflight has not passed" + control.execute_trajectory.assert_not_called() + + +def test_preflight_requires_configured_coordinator_task_without_loading_backend( + runtime: tuple[RuntimeModule, Any], +) -> None: + module, control = runtime + now = time.time() + _provide(module, top_ts=now, other_ts=now) + control.list_tasks.return_value = ["another_task"] + + status = module.preflight_rollout() + + assert status["policy_ready"] is False + assert "missing configured rollout task" in (status["last_error"] or "") + assert module._backend.load_count == 0 + control.execute_trajectory.assert_not_called() + + +def test_policy_actions_are_clipped_to_backend_bounds( + runtime: tuple[RuntimeModule, Any], +) -> None: + module, control = runtime + module._backend.actions[:, 1] = 2.0 + module._backend.info = PolicyBackendInfo( + name="fake", + chunk_length=30, + preferred_execution_steps=1, + action_lower=np.asarray([-100.0, 0.0], dtype=np.float32), + action_upper=np.asarray([100.0, 1.0], dtype=np.float32), + ) + now = time.time() + _provide(module, top_ts=now, other_ts=now) + + assert module.preflight_rollout()["policy_ready"] is True + module.start_rollout() + wait_until(lambda: control.execute_trajectory.call_count >= 1, timeout=1.0) + module.stop_rollout() + + trajectory = control.execute_trajectory.call_args.args[0] + assert trajectory.points[1].positions[1] == 1.0 + + +@pytest.mark.parametrize( + ("actions", "message"), + [ + (np.zeros((30, 1), dtype=np.float32), "expected (steps, 2)"), + (np.full((30, 2), np.nan, dtype=np.float32), "non-finite joint targets"), + ], +) +def test_invalid_action_chunk_cancels_and_latches_rollout_off( + runtime: tuple[RuntimeModule, Any], + actions: NDArray[np.float32], + message: str, +) -> None: + module, control = runtime + module._backend.actions = actions + now = time.time() + _provide(module, top_ts=now, other_ts=now) + + assert module.preflight_rollout()["policy_ready"] is True + module.start_rollout() + + wait_until(lambda: module.rollout_status()["active"] is False, timeout=1.0) + assert message in (module.rollout_status()["last_error"] or "") + control.execute_trajectory.assert_not_called() + control.cancel_trajectory.assert_called_with(task_name="policy_rollout") + + +def test_trajectory_rejection_cancels_and_latches_rollout_off( + runtime: tuple[RuntimeModule, Any], +) -> None: + module, control = runtime + now = time.time() + _provide(module, top_ts=now, other_ts=now) + control.execute_trajectory.return_value = TrajectoryExecutionResult( + TrajectoryExecutionStatus.POSITION_LIMIT_VIOLATION, + "outside hardware limits", + ) + + assert module.preflight_rollout()["policy_ready"] is True + module.start_rollout() + + wait_until(lambda: module.rollout_status()["active"] is False, timeout=1.0) + assert module.rollout_status()["last_error"] == "outside hardware limits" + control.cancel_trajectory.assert_called_with(task_name="policy_rollout") + + +def test_backend_reset_failure_does_not_leave_rollout_active( + runtime: tuple[RuntimeModule, Any], +) -> None: + module, control = runtime + now = time.time() + _provide(module, top_ts=now, other_ts=now) + assert module.preflight_rollout()["policy_ready"] is True + module._backend.reset_error = "backend state is stuck" + + module.start_rollout() + + wait_until(lambda: module.rollout_status()["active"] is False, timeout=1.0) + assert "backend state is stuck" in (module.rollout_status()["last_error"] or "") + control.cancel_trajectory.assert_called_with(task_name="policy_rollout") diff --git a/dimos/imitation/profile.py b/dimos/imitation/profile.py new file mode 100644 index 0000000000..2b2261d852 --- /dev/null +++ b/dimos/imitation/profile.py @@ -0,0 +1,159 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Typed observation and action contracts shared by collection and rollout.""" + +from __future__ import annotations + +import keyword + +from pydantic import Field, model_validator + +from dimos.imitation.dataprep.core import ( + DataPrepConfig, + FeatureSpec, + OutputConfig, + QualityConfig, + SyncConfig, +) +from dimos.protocol.service.spec import BaseConfig + + +class ImageSource(BaseConfig): + """One RGB image stream consumed under a backend feature key.""" + + stream: str + shape: tuple[int, int, int] + + @model_validator(mode="after") + def validate_source(self) -> ImageSource: + _validate_stream_name(self.stream) + if any(dimension <= 0 for dimension in self.shape): + raise ValueError("image shape must contain positive dimensions") + if self.shape[2] != 3: + raise ValueError("image source must be RGB HWC with three channels") + return self + + +class JointPositionSource(BaseConfig): + """One named JointState position projection.""" + + stream: str + joints: tuple[str, ...] = Field(min_length=1) + + @model_validator(mode="after") + def validate_source(self) -> JointPositionSource: + _validate_stream_name(self.stream) + if any(not joint.strip() for joint in self.joints): + raise ValueError("joint names must not be blank") + if len(set(self.joints)) != len(self.joints): + raise ValueError("joint names must be unique") + return self + + +PolicySource = ImageSource | JointPositionSource + + +class JointPositionAction(BaseConfig): + """Backend action key and the demonstration stream that teaches it.""" + + key: str = Field(min_length=1) + demonstration: JointPositionSource + + +class PolicyIOProfile(BaseConfig): + """Complete feature-key to DimOS-stream contract for one robot setup.""" + + name: str = Field(min_length=1) + robot_type: str = Field(min_length=1) + observations: dict[str, PolicySource] = Field(min_length=1) + action: JointPositionAction + sync: SyncConfig + quality: QualityConfig = QualityConfig() + + @model_validator(mode="after") + def validate_contract(self) -> PolicyIOProfile: + if any(not key.strip() for key in self.observations): + raise ValueError("observation feature keys must not be blank") + if self.action.key in self.observations: + raise ValueError("action key must not also be an observation key") + if self.sync.anchor not in self.observations: + raise ValueError("sync anchor must name an observation feature") + + types_by_stream: dict[str, type[PolicySource]] = {} + sources = [*self.observations.values(), self.action.demonstration] + for source in sources: + existing = types_by_stream.setdefault(source.stream, type(source)) + if existing is not type(source): + raise ValueError( + f"stream {source.stream!r} is declared with conflicting source types" + ) + matching_states = [ + key + for key, source in self.observations.items() + if isinstance(source, JointPositionSource) + and source.joints == self.action.demonstration.joints + ] + if len(matching_states) != 1: + raise ValueError( + "profile must have exactly one joint observation matching the action joints" + ) + return self + + @property + def action_state_key(self) -> str: + """Return the observation feature used as a trajectory's current state.""" + return next( + key + for key, source in self.observations.items() + if isinstance(source, JointPositionSource) + and source.joints == self.action.demonstration.joints + ) + + def dataprep_config(self, *, source: str = "", output: OutputConfig) -> DataPrepConfig: + """Project this live contract into the native-recording dataset schema.""" + observations = { + key: _feature_spec(source_spec) for key, source_spec in self.observations.items() + } + return DataPrepConfig( + source=source, + observation=observations, + action={self.action.key: _feature_spec(self.action.demonstration)}, + sync=self.sync, + quality=self.quality, + output=output, + ) + + +def _feature_spec(source: PolicySource) -> FeatureSpec: + if isinstance(source, ImageSource): + return FeatureSpec( + stream=source.stream, + field="data", + dtype="video", + shape=source.shape, + names=["height", "width", "channels"], + ) + return FeatureSpec( + stream=source.stream, + field="position", + dtype="float32", + shape=(len(source.joints),), + names=list(source.joints), + ) + + +def _validate_stream_name(name: str) -> None: + if not name.isidentifier() or keyword.iskeyword(name): + raise ValueError(f"stream {name!r} must be a Python identifier") diff --git a/dimos/imitation/test_profile.py b/dimos/imitation/test_profile.py new file mode 100644 index 0000000000..d5d6a7b330 --- /dev/null +++ b/dimos/imitation/test_profile.py @@ -0,0 +1,107 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from pathlib import Path + +from pydantic import ValidationError +import pytest + +from dimos.imitation.dataprep.core import OutputConfig, SyncConfig +from dimos.imitation.profile import ( + ImageSource, + JointPositionAction, + JointPositionSource, + PolicyIOProfile, +) + + +def _profile() -> PolicyIOProfile: + joints = ("left_joint", "right_joint") + return PolicyIOProfile( + name="dual-test", + robot_type="dual_test", + observations={ + "left": ImageSource(stream="left_image", shape=(480, 640, 3)), + "state": JointPositionSource(stream="joint_state", joints=joints), + }, + action=JointPositionAction( + key="actions", + demonstration=JointPositionSource(stream="joint_command", joints=joints), + ), + sync=SyncConfig(anchor="left", rate_hz=30.0, tolerance_ms=20.0), + ) + + +def test_profile_builds_matching_dataprep_config(tmp_path: Path) -> None: + profile = _profile() + output = OutputConfig(path=tmp_path / "dataset") + + config = profile.dataprep_config(source="recording.mcap", output=output) + + assert config.source == "recording.mcap" + assert config.output == output + assert config.observation["left"].stream == "left_image" + assert config.observation["left"].shape == (480, 640, 3) + assert config.observation["state"].names == ["left_joint", "right_joint"] + assert config.action["actions"].stream == "joint_command" + assert config.sync.anchor == "left" + + +@pytest.mark.parametrize( + ("update", "message"), + [ + ({"sync": SyncConfig(anchor="missing", rate_hz=30.0, tolerance_ms=20.0)}, "anchor"), + ( + { + "action": JointPositionAction( + key="left", + demonstration=JointPositionSource(stream="joint_command", joints=("joint",)), + ) + }, + "action key", + ), + ], +) +def test_profile_rejects_ambiguous_feature_contracts( + update: dict[str, object], message: str +) -> None: + values = _profile().model_dump() + values.update(update) + + with pytest.raises(ValidationError, match=message): + PolicyIOProfile.model_validate(values) + + +def test_profile_rejects_conflicting_types_on_one_stream() -> None: + values = _profile().model_dump() + values["action"] = { + "key": "actions", + "demonstration": {"stream": "left_image", "joints": ["joint"]}, + } + + with pytest.raises(ValidationError, match="conflicting source types"): + PolicyIOProfile.model_validate(values) + + +@pytest.mark.parametrize( + "source", + [ + {"stream": "left-image", "shape": (480, 640, 3)}, + {"stream": "left_image", "shape": (480, 640, 1)}, + {"stream": "left_image", "shape": (0, 640, 3)}, + ], +) +def test_image_source_requires_an_rgb_hwc_python_port(source: dict[str, object]) -> None: + with pytest.raises(ValidationError): + ImageSource.model_validate(source) diff --git a/dimos/robot/all_blueprints.py b/dimos/robot/all_blueprints.py index f62592d68e..bb9cd4b0d3 100644 --- a/dimos/robot/all_blueprints.py +++ b/dimos/robot/all_blueprints.py @@ -255,7 +255,6 @@ "joystick-module": "dimos.robot.unitree.b1.joystick_module.JoystickModule", "keyboard-teleop": "dimos.robot.unitree.keyboard_teleop.KeyboardTeleop", "keyboard-teleop-module": "dimos.teleop.keyboard.keyboard_teleop_module.KeyboardTeleopModule", - "le-robot-policy-module": "dimos.imitation.policy.lerobot.module.LeRobotPolicyModule", "lidar-window-relocalization": "dimos.mapping.relocalization.lidar.module.LidarWindowRelocalization", "local-map-relocalization": "dimos.mapping.relocalization.lidar.module.LocalMapRelocalization", "m20-camera-relay": "dimos.robot.deeprobotics.m20.camera.M20CameraRelay", diff --git a/dimos/robot/manipulators/dual_openyam/blueprints/learning_collection.py b/dimos/robot/manipulators/dual_openyam/blueprints/learning_collection.py new file mode 100644 index 0000000000..57b7806750 --- /dev/null +++ b/dimos/robot/manipulators/dual_openyam/blueprints/learning_collection.py @@ -0,0 +1,56 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Dual OpenYAM Quest collection with two independently declared cameras.""" + +from pathlib import Path + +from dimos.core.coordination.blueprints import Blueprint, autoconnect +from dimos.experimental.memory.rust_recorder import RustMcapStoreConfig +from dimos.imitation.cameras import CameraDevice, profile_cameras +from dimos.imitation.collection.episode_monitor import EpisodeMonitorModule +from dimos.robot.manipulators.dual_openyam.blueprints.teleop import ( + build_dual_openyam_quest_teleop, +) +from dimos.robot.manipulators.dual_openyam.learning import ( + DUAL_OPENYAM_TWO_WRIST_IO, + DualOpenYamQuestRecorder, +) + + +def build_dual_openyam_quest_collection( + *, + recording: Path, + task: str, + cameras: dict[str, CameraDevice], + left_can_port: str | None = None, + right_can_port: str | None = None, +) -> Blueprint: + """Build a bimanual Quest collection session and two wrist cameras.""" + camera_blueprints, camera_remappings = profile_cameras( + DUAL_OPENYAM_TWO_WRIST_IO, + cameras, + ) + return autoconnect( + DualOpenYamQuestRecorder.blueprint( + store=RustMcapStoreConfig(path=str(recording)), + record_tf=False, + ), + EpisodeMonitorModule.blueprint(task=task), + build_dual_openyam_quest_teleop( + left_can_port=left_can_port, + right_can_port=right_can_port, + ), + *camera_blueprints, + ).remappings(camera_remappings) diff --git a/dimos/robot/manipulators/dual_openyam/blueprints/learning_rollout.py b/dimos/robot/manipulators/dual_openyam/blueprints/learning_rollout.py new file mode 100644 index 0000000000..0d2a40fc3d --- /dev/null +++ b/dimos/robot/manipulators/dual_openyam/blueprints/learning_rollout.py @@ -0,0 +1,78 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Released ABC-DiT rollout for the complete Dual OpenYAM entity.""" + +from dimos.constants import DEFAULT_CAPACITY_COLOR_IMAGE +from dimos.control.coordinator import TaskConfig +from dimos.core.coordination.blueprints import Blueprint, autoconnect +from dimos.core.transport import pSHMTransport +from dimos.imitation.cameras import CameraDevice, profile_cameras +from dimos.imitation.policy.abc.module import DualOpenYamAbcPolicy +from dimos.imitation.policy.module import ( + POLICY_ROLLOUT_INSTANCE_NAME, + POLICY_ROLLOUT_TASK_NAME, +) +from dimos.msgs.sensor_msgs.Image import Image +from dimos.robot.manipulators.dual_openyam.blueprints.basic import DualOpenYamCoordinator +from dimos.robot.manipulators.dual_openyam.learning import ABC_JOINTS, DUAL_OPENYAM_ABC_IO + + +def build_dual_openyam_abc_rollout( + *, + artifact: str, + task: str, + cameras: dict[str, CameraDevice], + device: str | None = None, + quest_control: bool = False, + left_can_port: str | None = None, + right_can_port: str | None = None, +) -> Blueprint: + """Build three-camera ABC rollout; a synthetic top view is never accepted.""" + if quest_control: + raise ValueError("Quest takeover is not defined for Dual OpenYAM ABC rollout") + camera_blueprints, camera_remappings = profile_cameras(DUAL_OPENYAM_ABC_IO, cameras) + blueprint = autoconnect( + DualOpenYamAbcPolicy.blueprint( + instance_name=POLICY_ROLLOUT_INSTANCE_NAME, + artifact=artifact, + task=task, + device=device, + trajectory_task_name=POLICY_ROLLOUT_TASK_NAME, + ), + DualOpenYamCoordinator.blueprint( + instance_name="ControlCoordinator", + left_can_port=left_can_port, + right_can_port=right_can_port, + tasks=[ + TaskConfig( + name=POLICY_ROLLOUT_TASK_NAME, + type="trajectory", + joint_names=list(ABC_JOINTS), + priority=10, + params={"start_position_tolerance": 0.05}, + ) + ], + ), + *camera_blueprints, + ).remappings(camera_remappings) + return blueprint.transports( + { + (stream, Image): pSHMTransport.spec( + f"/{stream}", + default_capacity=DEFAULT_CAPACITY_COLOR_IMAGE, + ) + for stream in ("top_image", "left_wrist_image", "right_wrist_image") + } + ) diff --git a/dimos/robot/manipulators/dual_openyam/blueprints/teleop.py b/dimos/robot/manipulators/dual_openyam/blueprints/teleop.py index 47727c4d07..1df57ff9fc 100644 --- a/dimos/robot/manipulators/dual_openyam/blueprints/teleop.py +++ b/dimos/robot/manipulators/dual_openyam/blueprints/teleop.py @@ -15,7 +15,7 @@ """Coupled Quest teleoperation for the complete Dual OpenYAM entity.""" from dimos.control.coordinator import TaskConfig -from dimos.core.coordination.blueprints import autoconnect +from dimos.core.coordination.blueprints import Blueprint, autoconnect from dimos.manipulation.manipulation_module import ManipulationModule from dimos.manipulation.planning.kinematics.config import PinkKinematicsConfig from dimos.robot.manipulators.common.blueprints import teleop_ik_task @@ -73,39 +73,51 @@ }, ) -teleop_quest_dual_openyam = autoconnect( - ArmTeleopModule.blueprint(), - DualOpenYamCoordinator.blueprint( - instance_name="ControlCoordinator", - tasks=[ - _dual_openyam_quest_task, - TaskConfig( - name="left_arm_gripper", - type="gripper", - joint_names=[DUAL_OPENYAM_GRIPPER_JOINTS[0]], - priority=20, - stream_bind={"gripper_command": "left_gripper_command"}, - ), - TaskConfig( - name="right_arm_gripper", - type="gripper", - joint_names=[DUAL_OPENYAM_GRIPPER_JOINTS[1]], - priority=20, - stream_bind={"gripper_command": "right_gripper_command"}, - ), - dual_openyam_trajectory_task(priority=20), - ], - ), - ManipulationModule.blueprint( - model=_dual_openyam_quest_model, - kinematics=_dual_openyam_quest_pink, - visualization={"backend": "viser"}, - ), -).remappings( - [ - (ArmTeleopModule, "left_controller_output", "left_cartesian_command"), - (ArmTeleopModule, "left_gripper_command", "left_gripper_command"), - (ArmTeleopModule, "right_controller_output", "right_cartesian_command"), - (ArmTeleopModule, "right_gripper_command", "right_gripper_command"), - ] -) + +def build_dual_openyam_quest_teleop( + *, + left_can_port: str | None = None, + right_can_port: str | None = None, +) -> Blueprint: + """Build Quest teleop against mock or explicitly selected dual-CAN hardware.""" + return autoconnect( + ArmTeleopModule.blueprint(), + DualOpenYamCoordinator.blueprint( + instance_name="ControlCoordinator", + left_can_port=left_can_port, + right_can_port=right_can_port, + tasks=[ + _dual_openyam_quest_task, + TaskConfig( + name="left_arm_gripper", + type="gripper", + joint_names=[DUAL_OPENYAM_GRIPPER_JOINTS[0]], + priority=20, + stream_bind={"gripper_command": "left_gripper_command"}, + ), + TaskConfig( + name="right_arm_gripper", + type="gripper", + joint_names=[DUAL_OPENYAM_GRIPPER_JOINTS[1]], + priority=20, + stream_bind={"gripper_command": "right_gripper_command"}, + ), + dual_openyam_trajectory_task(priority=20), + ], + ), + ManipulationModule.blueprint( + model=_dual_openyam_quest_model, + kinematics=_dual_openyam_quest_pink, + visualization={"backend": "viser"}, + ), + ).remappings( + [ + (ArmTeleopModule, "left_controller_output", "left_cartesian_command"), + (ArmTeleopModule, "left_gripper_command", "left_gripper_command"), + (ArmTeleopModule, "right_controller_output", "right_cartesian_command"), + (ArmTeleopModule, "right_gripper_command", "right_gripper_command"), + ] + ) + + +teleop_quest_dual_openyam = autoconnect(build_dual_openyam_quest_teleop()) diff --git a/dimos/robot/manipulators/dual_openyam/blueprints/test_learning.py b/dimos/robot/manipulators/dual_openyam/blueprints/test_learning.py new file mode 100644 index 0000000000..6fd93478e9 --- /dev/null +++ b/dimos/robot/manipulators/dual_openyam/blueprints/test_learning.py @@ -0,0 +1,86 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from pathlib import Path + +import pytest + +from dimos.hardware.sensors.camera.module import CameraModule +from dimos.imitation.policy.abc.module import DualOpenYamAbcPolicy +from dimos.imitation.policy.module import POLICY_ROLLOUT_INSTANCE_NAME, POLICY_ROLLOUT_TASK_NAME +from dimos.robot.manipulators.dual_openyam.blueprints.learning_collection import ( + build_dual_openyam_quest_collection, +) +from dimos.robot.manipulators.dual_openyam.blueprints.learning_rollout import ( + build_dual_openyam_abc_rollout, +) +from dimos.robot.manipulators.dual_openyam.learning import ABC_JOINTS, DualOpenYamQuestRecorder + + +def test_dual_collection_declares_two_distinct_camera_sources(tmp_path: Path) -> None: + blueprint = build_dual_openyam_quest_collection( + recording=tmp_path / "dual.mcap", + task="fold towel", + cameras={"left_wrist_image": 0, "right_wrist_image": 1}, + left_can_port="follower_l", + right_can_port="follower_r", + ) + cameras = [atom for atom in blueprint.active_blueprints if atom.module is CameraModule] + + assert any(atom.module is DualOpenYamQuestRecorder for atom in blueprint.active_blueprints) + assert [camera.kwargs["hardware"].camera_index for camera in cameras] == [0, 1] + assert blueprint.remapping_map[("PolicyCamera_left_wrist_image", "color_image")] == ( + "left_wrist_image" + ) + assert blueprint.remapping_map[("PolicyCamera_right_wrist_image", "color_image")] == ( + "right_wrist_image" + ) + coordinator = next( + atom for atom in blueprint.active_blueprints if atom.name == "ControlCoordinator" + ) + assert coordinator.kwargs["left_can_port"] == "follower_l" + assert coordinator.kwargs["right_can_port"] == "follower_r" + + +def test_abc_rollout_requires_all_three_physical_cameras() -> None: + with pytest.raises(ValueError, match="top_image"): + build_dual_openyam_abc_rollout( + artifact="checkpoint.pt", + task="put bottles in bin", + cameras={"left_wrist_image": 0, "right_wrist_image": 1}, + ) + + +def test_abc_rollout_uses_released_action_order_and_stable_rpc_name() -> None: + blueprint = build_dual_openyam_abc_rollout( + artifact="checkpoint.pt", + task="put bottles in bin", + cameras={"top_image": 0, "left_wrist_image": 1, "right_wrist_image": 2}, + left_can_port="follower_l", + right_can_port="follower_r", + ) + policy = next( + atom for atom in blueprint.active_blueprints if atom.module is DualOpenYamAbcPolicy + ) + coordinator = next( + atom for atom in blueprint.active_blueprints if atom.name == "ControlCoordinator" + ) + + assert policy.name == POLICY_ROLLOUT_INSTANCE_NAME + rollout_task = next( + task for task in coordinator.kwargs["tasks"] if task.name == POLICY_ROLLOUT_TASK_NAME + ) + assert rollout_task.joint_names == list(ABC_JOINTS) + assert coordinator.kwargs["left_can_port"] == "follower_l" + assert coordinator.kwargs["right_can_port"] == "follower_r" diff --git a/dimos/robot/manipulators/dual_openyam/learning.py b/dimos/robot/manipulators/dual_openyam/learning.py new file mode 100644 index 0000000000..01b8852575 --- /dev/null +++ b/dimos/robot/manipulators/dual_openyam/learning.py @@ -0,0 +1,107 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Distinct Dual OpenYAM collection and released-ABC rollout profiles.""" + +from dimos.imitation.collection.native_recorder import declare_recorder +from dimos.imitation.dataprep.core import QualityConfig, SyncConfig +from dimos.imitation.profile import ( + ImageSource, + JointPositionAction, + JointPositionSource, + PolicyIOProfile, +) +from dimos.robot.manipulators.dual_openyam.config import ( + DUAL_OPENYAM_GRIPPER_JOINTS, + DUAL_OPENYAM_JOINTS, + DUAL_OPENYAM_LEFT_ARM_JOINTS, + DUAL_OPENYAM_RIGHT_ARM_JOINTS, +) + +DUAL_OPENYAM_CAMERA_SHAPE = (480, 640, 3) +DUAL_OPENYAM_FPS = 30.0 +ABC_JOINTS = ( + *DUAL_OPENYAM_LEFT_ARM_JOINTS, + DUAL_OPENYAM_GRIPPER_JOINTS[0], + *DUAL_OPENYAM_RIGHT_ARM_JOINTS, + DUAL_OPENYAM_GRIPPER_JOINTS[1], +) + +_quality = QualityConfig( + mode="strict", + min_source_rate_ratio=0.95, + max_camera_gap_ms=100.0, + max_alignment_error_ms=20.0, +) + +DUAL_OPENYAM_TWO_WRIST_IO = PolicyIOProfile( + name="dual-openyam-quest", + robot_type="dual_openyam", + observations={ + "observation.images.left_wrist": ImageSource( + stream="left_wrist_image", + shape=DUAL_OPENYAM_CAMERA_SHAPE, + ), + "observation.images.right_wrist": ImageSource( + stream="right_wrist_image", + shape=DUAL_OPENYAM_CAMERA_SHAPE, + ), + "observation.state": JointPositionSource( + stream="coordinator_joint_state", + joints=tuple(DUAL_OPENYAM_JOINTS), + ), + }, + action=JointPositionAction( + key="action", + demonstration=JointPositionSource( + stream="applied_joint_position_command", + joints=tuple(DUAL_OPENYAM_JOINTS), + ), + ), + sync=SyncConfig( + anchor="observation.images.left_wrist", + rate_hz=DUAL_OPENYAM_FPS, + tolerance_ms=20.0, + ), + quality=_quality, +) + +DUAL_OPENYAM_ABC_IO = PolicyIOProfile( + name="dual-openyam-abc", + robot_type="dual_openyam", + observations={ + "top": ImageSource(stream="top_image", shape=DUAL_OPENYAM_CAMERA_SHAPE), + "left": ImageSource(stream="left_wrist_image", shape=DUAL_OPENYAM_CAMERA_SHAPE), + "right": ImageSource(stream="right_wrist_image", shape=DUAL_OPENYAM_CAMERA_SHAPE), + "state": JointPositionSource( + stream="coordinator_joint_state", + joints=ABC_JOINTS, + ), + }, + action=JointPositionAction( + key="actions", + demonstration=JointPositionSource( + stream="applied_joint_position_command", + joints=ABC_JOINTS, + ), + ), + sync=SyncConfig(anchor="top", rate_hz=DUAL_OPENYAM_FPS, tolerance_ms=20.0), + quality=_quality, +) + +DualOpenYamQuestRecorder = declare_recorder( + "DualOpenYamQuestRecorder", + __name__, + DUAL_OPENYAM_TWO_WRIST_IO, +) diff --git a/docs/capabilities/manipulation/imitation-learning.md b/docs/capabilities/manipulation/imitation-learning.md index 1e1e25968d..05acd36fa9 100644 --- a/docs/capabilities/manipulation/imitation-learning.md +++ b/docs/capabilities/manipulation/imitation-learning.md @@ -9,13 +9,17 @@ TUI attaches to its episode-control interface; it does not own the robot. | --- | --- | --- | | `openyam-teach-collection` | Wrist RGB | Measured 7-D joints for both | | `openyam-quest-collection` | Wrist RGB | Measured state, accepted commands | +| `dual-openyam-quest-collection` | Two wrist RGB cameras | Measured state, accepted commands; 14-D | ```bash -dimos --can-port follower_l run openyam-teach-collection --daemon \ - --recorder.recording recordings/session-001 \ +dimos run dual-openyam-quest-collection --daemon \ + --recorder.recording recordings/fold-001 \ --recorder.format mcap \ - --episodes.task "pick up the cube" \ - --wrist.hardware.camera-index /dev/video0 + --episodes.task "fold the towel" \ + --controlcoordinator.left-can-port follower_l \ + --controlcoordinator.right-can-port follower_r \ + --left-wrist.hardware.camera-index /dev/video0 \ + --right-wrist.hardware.camera-index /dev/video2 dimos imitation collect ``` @@ -29,10 +33,14 @@ episode, Q asks for confirmation: **recording and the robot continue after the TUI exits**. Use `dimos stop` separately to stop the stack; stopping real hardware may de-torque the arms, so support them first. +The dual profile uses 640×480 RGB images at 30 Hz with a 20 ms alignment +tolerance anchored on the left wrist. Joint order is left arm joints 1–6, +right arm joints 1–6, left gripper, right gripper. + ## Recording directories ```text -recordings/session-001/ +recordings/fold-001/ ├── schema.json └── recording.mcap ``` @@ -105,13 +113,13 @@ The controller class must be importable in the client. ## Prepare and train ```bash -dimos imitation inspect recordings/session-001 -dimos imitation prepare recordings/session-001 --output datasets/session-001 +dimos imitation inspect recordings/fold-001 +dimos imitation prepare recordings/fold-001 --output datasets/fold dimos imitation train \ - --dataset.repo_id=local/openyam-teach \ - --dataset.root=datasets/session-001 \ + --dataset.repo_id=local/dual-openyam-quest \ + --dataset.root=datasets/fold \ --policy.type=act \ - --output_dir=outputs/openyam-act + --output_dir=outputs/dual-openyam-act ``` Preparation reads the saved schema, not the current robot blueprint. Python diff --git a/pyproject.toml b/pyproject.toml index 5c6af291ce..3e36ba43bc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -54,6 +54,9 @@ exclude = [ # shipping arbitrary manifests from the source tree. "imitation/policy/lerobot/python/pyproject.toml", "imitation/policy/lerobot/python/uv.lock", + "imitation/policy/abc/python/pyproject.toml", + "imitation/policy/abc/python/uv.lock", + "imitation/policy/abc/python/VENDORED.md", ] [tool.setuptools.exclude-package-data] @@ -595,6 +598,7 @@ exclude = [ "venv", "libs", "external", + "dimos/imitation/policy/abc/python/abc_minimal", "src" ] @@ -638,7 +642,7 @@ strict = true warn_unused_ignores = false untyped_calls_exclude = ["zenoh"] explicit_package_bases = true -exclude = "^dimos/models/Detic(/|$)|^dimos/imitation/policy/lerobot/python/|.*/test_.|.*/tool_.|.*/conftest.py*" +exclude = "^dimos/models/Detic(/|$)|^dimos/imitation/policy/(lerobot|abc)/python/|.*/test_.|.*/tool_.|.*/conftest.py*" [[tool.mypy.overrides]] module = [ From ebf3fb853fa31b5e4d401ce9ac85ea32bf7d0854 Mon Sep 17 00:00:00 2001 From: cc Date: Wed, 9 Sep 2026 15:23:54 -0700 Subject: [PATCH 2/4] refactor(imitation): scope profiles to dual-arm collection --- .pre-commit-config.yaml | 2 +- MANIFEST.in | 11 - dimos/imitation/cameras.py | 18 +- dimos/imitation/dataprep/lerobot.py | 4 +- dimos/imitation/policy/abc/module.py | 47 - dimos/imitation/policy/abc/python/VENDORED.md | 11 - .../policy/abc/python/abc_minimal/__init__.py | 1 - .../policy/abc/python/abc_minimal/config.py | 48 - .../policy/abc/python/abc_minimal/dit.py | 895 ------------------ .../abc/python/abc_minimal/fast_inference.py | 118 --- .../abc/python/abc_minimal/preprocess.py | 71 -- .../policy/abc/python/dimos_abc/__init__.py | 1 - .../policy/abc/python/dimos_abc/runtime.py | 208 ---- .../abc/python/dimos_abc/runtime_tests.py | 40 - .../policy/abc/python/pyproject.toml | 56 -- dimos/imitation/policy/abc/python/uv.lock | 569 ----------- dimos/imitation/policy/abc/test_module.py | 45 - dimos/imitation/policy/backend.py | 53 -- dimos/imitation/policy/lerobot/README.md | 63 +- .../lerobot/python/dimos_lerobot/runtime.py | 548 +++++++++-- .../python/dimos_lerobot/runtime_tests.py | 655 ++++++++++++- .../policy/lerobot/python/pyproject.toml | 4 - dimos/imitation/policy/lerobot/test_module.py | 69 +- dimos/imitation/policy/module.py | 161 ---- dimos/imitation/policy/runtime.py | 485 ---------- dimos/imitation/policy/test_runtime.py | 320 ------- dimos/imitation/profile.py | 159 ---- dimos/imitation/test_profile.py | 107 --- dimos/robot/all_blueprints.py | 1 + .../blueprints/learning_collection.py | 31 +- .../blueprints/learning_rollout.py | 78 -- .../dual_openyam/blueprints/test_learning.py | 86 +- .../robot/manipulators/dual_openyam/config.py | 21 +- .../robot/manipulators/dual_openyam/joints.py | 22 + .../manipulators/dual_openyam/learning.py | 112 +-- pyproject.toml | 6 +- 36 files changed, 1309 insertions(+), 3817 deletions(-) delete mode 100644 dimos/imitation/policy/abc/module.py delete mode 100644 dimos/imitation/policy/abc/python/VENDORED.md delete mode 100644 dimos/imitation/policy/abc/python/abc_minimal/__init__.py delete mode 100644 dimos/imitation/policy/abc/python/abc_minimal/config.py delete mode 100644 dimos/imitation/policy/abc/python/abc_minimal/dit.py delete mode 100644 dimos/imitation/policy/abc/python/abc_minimal/fast_inference.py delete mode 100644 dimos/imitation/policy/abc/python/abc_minimal/preprocess.py delete mode 100644 dimos/imitation/policy/abc/python/dimos_abc/__init__.py delete mode 100644 dimos/imitation/policy/abc/python/dimos_abc/runtime.py delete mode 100644 dimos/imitation/policy/abc/python/dimos_abc/runtime_tests.py delete mode 100644 dimos/imitation/policy/abc/python/pyproject.toml delete mode 100644 dimos/imitation/policy/abc/python/uv.lock delete mode 100644 dimos/imitation/policy/abc/test_module.py delete mode 100644 dimos/imitation/policy/backend.py delete mode 100644 dimos/imitation/policy/module.py delete mode 100644 dimos/imitation/policy/runtime.py delete mode 100644 dimos/imitation/policy/test_runtime.py delete mode 100644 dimos/imitation/profile.py delete mode 100644 dimos/imitation/test_profile.py delete mode 100644 dimos/robot/manipulators/dual_openyam/blueprints/learning_rollout.py create mode 100644 dimos/robot/manipulators/dual_openyam/joints.py diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 2825b53048..de6ab3e28f 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,6 +1,6 @@ default_stages: [pre-commit] default_install_hook_types: [pre-commit, commit-msg] -exclude: (dimos/models/.*)|(deprecated)|(dimos/imitation/policy/abc/python/abc_minimal/) +exclude: (dimos/models/.*)|(deprecated) repos: - repo: https://github.com/Lucas-C/pre-commit-hooks diff --git a/MANIFEST.in b/MANIFEST.in index e9e13ae147..912e9ad0a9 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -11,9 +11,6 @@ global-exclude .DS_Store recursive-include dimos *.yaml *.yml *.json *.urdf *.html *.css *.js *.svg *.tcss include dimos/imitation/policy/lerobot/python/pyproject.toml include dimos/imitation/policy/lerobot/python/uv.lock -include dimos/imitation/policy/abc/python/pyproject.toml -include dimos/imitation/policy/abc/python/uv.lock -include dimos/imitation/policy/abc/python/VENDORED.md # --- Exclusions (must come after the includes above so they win) --- # Test fixtures must never ship. @@ -48,14 +45,6 @@ prune .pytest_cache prune .ruff_cache prune .vscode prune dimos/web/command-center-extension -prune dimos/imitation/policy/lerobot/python/.mypy_cache -prune dimos/imitation/policy/lerobot/python/.pytest_cache -prune dimos/imitation/policy/lerobot/python/.ruff_cache -prune dimos/imitation/policy/lerobot/python/.venv -prune dimos/imitation/policy/abc/python/.mypy_cache -prune dimos/imitation/policy/abc/python/.pytest_cache -prune dimos/imitation/policy/abc/python/.ruff_cache -prune dimos/imitation/policy/abc/python/.venv global-exclude test_*.py global-exclude conftest.py diff --git a/dimos/imitation/cameras.py b/dimos/imitation/cameras.py index 153c078d2d..df63be25aa 100644 --- a/dimos/imitation/cameras.py +++ b/dimos/imitation/cameras.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Camera blueprints generated from policy image-source declarations.""" +"""Camera blueprints generated from collection image features.""" from __future__ import annotations @@ -21,21 +21,17 @@ from dimos.core.coordination.blueprints import Blueprint from dimos.hardware.sensors.camera.module import CameraModule from dimos.hardware.sensors.camera.webcam import WebcamConfig -from dimos.imitation.profile import ImageSource, PolicyIOProfile +from dimos.imitation.collection.profile import CollectionProfile CameraDevice = int | str def profile_cameras( - profile: PolicyIOProfile, + profile: CollectionProfile, devices: Mapping[str, CameraDevice], ) -> tuple[list[Blueprint], list[tuple[str, str, str]]]: - """Build cameras and explicit output remappings for a policy profile.""" - image_sources = { - source.stream: source - for source in profile.observations.values() - if isinstance(source, ImageSource) - } + """Build cameras and explicit output remappings for a collection profile.""" + image_sources = profile.camera_features() missing = sorted(set(image_sources) - set(devices)) unknown = sorted(set(devices) - set(image_sources)) if missing or unknown: @@ -49,8 +45,10 @@ def profile_cameras( blueprints: list[Blueprint] = [] remappings: list[tuple[str, str, str]] = [] for stream_name, source in image_sources.items(): + if len(source.shape) != 3 or source.shape[2] != 3: + raise ValueError(f"Camera {stream_name!r} requires an HWC RGB shape") height, width, _channels = source.shape - instance_name = f"PolicyCamera_{stream_name}" + instance_name = f"CollectionCamera_{stream_name}" blueprints.append( CameraModule.blueprint( instance_name=instance_name, diff --git a/dimos/imitation/dataprep/lerobot.py b/dimos/imitation/dataprep/lerobot.py index 0422b709df..b6fc404330 100644 --- a/dimos/imitation/dataprep/lerobot.py +++ b/dimos/imitation/dataprep/lerobot.py @@ -33,13 +33,13 @@ Result, ) from dimos.imitation.dataprep.core import DataPrepConfig -from dimos.imitation.policy.lerobot.module import OpenYamLeRobotPolicy +from dimos.imitation.policy.lerobot.module import LeRobotPolicyModule from dimos.utils.cache import cache_usage_guard def lerobot_project() -> Path: """Locate the packaged LeRobot project beside its host contract.""" - source = Path(inspect.getfile(OpenYamLeRobotPolicy)).resolve() + source = Path(inspect.getfile(LeRobotPolicyModule)).resolve() return source.parent / "python" diff --git a/dimos/imitation/policy/abc/module.py b/dimos/imitation/policy/abc/module.py deleted file mode 100644 index 0c58991194..0000000000 --- a/dimos/imitation/policy/abc/module.py +++ /dev/null @@ -1,47 +0,0 @@ -# Copyright 2026 Dimensional Inc. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Dual OpenYAM binding for Amazon's released ABC-DiT checkpoint.""" - -from pathlib import Path - -from pydantic import Field, field_validator - -from dimos.imitation.policy.module import PolicyRolloutConfig, declare_policy_module -from dimos.robot.manipulators.dual_openyam.learning import DUAL_OPENYAM_ABC_IO - - -class AbcPolicyConfig(PolicyRolloutConfig): - """Released ABC-DiT inference settings.""" - - norm_stats_path: str | None = None - diffusion_steps: int = Field(default=10, ge=1) - fast_inference: bool = True - - @field_validator("norm_stats_path") - @classmethod - def resolve_norm_stats_path(cls, value: str | None) -> str | None: - if value is None: - return None - path = Path(value).expanduser() - return str(path.resolve()) if path.exists() else value - - -DualOpenYamAbcPolicy = declare_policy_module( - "DualOpenYamAbcPolicy", - __name__, - DUAL_OPENYAM_ABC_IO, - AbcPolicyConfig, - "dimos_abc.runtime:AbcPolicyRuntime", -) diff --git a/dimos/imitation/policy/abc/python/VENDORED.md b/dimos/imitation/policy/abc/python/VENDORED.md deleted file mode 100644 index 0249a1f13b..0000000000 --- a/dimos/imitation/policy/abc/python/VENDORED.md +++ /dev/null @@ -1,11 +0,0 @@ -# Vendored ABC inference code - -`abc_minimal/dit.py`, `preprocess.py`, and `fast_inference.py` come from -[`amazon-far/abc`](https://github.com/amazon-far/abc) at revision -`6bc6586721cf0c409ccee80f675a28de9b9b2f5e`. `config.py` retains only the two -configuration dataclasses used for inference, and `fast_inference.py` omits the -RTC helper. The upstream project is licensed under Apache-2.0, the same license -as this repository. - -Training, simulation, dataset conversion, visualization, and RTC code are not -vendored. diff --git a/dimos/imitation/policy/abc/python/abc_minimal/__init__.py b/dimos/imitation/policy/abc/python/abc_minimal/__init__.py deleted file mode 100644 index 2bcd1fedfc..0000000000 --- a/dimos/imitation/policy/abc/python/abc_minimal/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Vendored ABC-DiT inference subset.""" diff --git a/dimos/imitation/policy/abc/python/abc_minimal/config.py b/dimos/imitation/policy/abc/python/abc_minimal/config.py deleted file mode 100644 index a31d826d4c..0000000000 --- a/dimos/imitation/policy/abc/python/abc_minimal/config.py +++ /dev/null @@ -1,48 +0,0 @@ -"""Inference configuration extracted from the upstream ABC minimal release.""" - -import os -from dataclasses import dataclass, field -from pathlib import Path - - -REPO_ROOT = Path(__file__).resolve().parents[1] -DEFAULT_CACHE_ROOT = REPO_ROOT / "cache" - - -def default_cache_root() -> Path: - return Path(os.environ.get("ABC_CACHE", str(DEFAULT_CACHE_ROOT))).expanduser() - - -@dataclass -class ClipConfig: - """CLIP ViT-B/32 text asset locations.""" - - cache_dir: str = field(default_factory=lambda: str(Path.home() / ".cache" / "clip")) - model_url: str = ( - "https://openaipublic.azureedge.net/clip/models/" - "40d365715913c9da98579312b702a82c18be219cc2a73407c4526f58eba950af/ViT-B-32.pt" - ) - bpe_url: str = "https://github.com/openai/CLIP/raw/main/clip/bpe_simple_vocab_16e6.txt.gz" - model_name: str = "ViT-B-32.pt" - bpe_name: str = "bpe_simple_vocab_16e6.txt.gz" - - -@dataclass -class DiTConfig: - """ABC-DiT architecture used by the released checkpoint.""" - - hidden_size: int = 1536 - depth: int = 32 - num_heads: int = 24 - mlp_ratio: float = 4.0 - state_dim: int = 14 - action_dim: int = 14 - chunk_length: int = 30 - camera_keys: tuple[str, ...] = ("top", "left", "right") - task_embed_dim: int = 512 - vit_embed_dim: int = 768 - vit_depth: int = 12 - vit_num_heads: int = 12 - vision_pool_num_queries: int = 12 - vision_pool_num_heads: int = 8 - vision_pool_mlp_ratio: int = 4 diff --git a/dimos/imitation/policy/abc/python/abc_minimal/dit.py b/dimos/imitation/policy/abc/python/abc_minimal/dit.py deleted file mode 100644 index 82c03dde00..0000000000 --- a/dimos/imitation/policy/abc/python/abc_minimal/dit.py +++ /dev/null @@ -1,895 +0,0 @@ -"""ABC-DiT policy implementation for the released bottles-in-bin checkpoints. - -Includes the CLIP text encoder and DINOv3 vision backbone needed to run the model. -""" - -import gzip -import html -import math -import urllib.request -from functools import lru_cache -from pathlib import Path - -import numpy as np -import torch -import torch.nn as nn -import torch.nn.functional as F - -from abc_minimal.config import ClipConfig, DiTConfig - -# CLIP ViT-B/32 text encoder. - -SOT_TOKEN = "<|startoftext|>" -EOT_TOKEN = "<|endoftext|>" - - -def task_name_to_prompt(task_name): - """Convert production task names like open_the_pen_caps to CLIP prompt text.""" - return " ".join(task_name.replace("-", " ").replace("_", " ").split()) - - -def _load_clip_text_deps(): - try: - import ftfy - import regex - except ImportError as exc: - raise RuntimeError( - "CLIP text embedding requires the 'ftfy' and 'regex' packages. " - "Run `uv sync` after pulling this version, or install them manually." - ) from exc - return ftfy, regex - - -def _download_if_missing(url, path): - path.parent.mkdir(parents=True, exist_ok=True) - if not path.exists(): - urllib.request.urlretrieve(url, path) - - -def ensure_clip_text_assets(config: ClipConfig): - """Download CLIP ViT-B/32 text assets if needed.""" - root = Path(config.cache_dir).expanduser() - b32_path = root / config.model_name - bpe_path = root / config.bpe_name - _download_if_missing(config.model_url, b32_path) - _download_if_missing(config.bpe_url, bpe_path) - return b32_path, bpe_path - - -@lru_cache() -def _bytes_to_unicode(): - bs = ( - list(range(ord("!"), ord("~") + 1)) - + list(range(ord("¡"), ord("¬") + 1)) - + list(range(ord("®"), ord("ÿ") + 1)) - ) - cs = bs[:] - n = 0 - for b in range(256): - if b not in bs: - bs.append(b) - cs.append(256 + n) - n += 1 - return dict(zip(bs, [chr(c) for c in cs])) - - -def _get_pairs(word): - return set(zip(word, word[1:])) - - -class CLIPBPETokenizer: - """Small copy of OpenAI CLIP's BPE tokenizer, scoped to text encoding.""" - - def __init__(self, bpe_path): - _ftfy, regex = _load_clip_text_deps() - self.ftfy = _ftfy - self.regex = regex - with gzip.open(bpe_path, "rt", encoding="utf-8") as f: - merges = [tuple(l.split()) for l in f.read().split("\n")[1 : 49152 - 256 - 2 + 1]] - self.byte_encoder = _bytes_to_unicode() - vocab = list(self.byte_encoder.values()) - vocab = vocab + [v + "" for v in vocab] - for merge in merges: - vocab.append("".join(merge)) - vocab.extend([SOT_TOKEN, EOT_TOKEN]) - self.encoder = {v: i for i, v in enumerate(vocab)} - self.bpe_ranks = dict(zip(merges, range(len(merges)))) - self.cache = {SOT_TOKEN: SOT_TOKEN, EOT_TOKEN: EOT_TOKEN} - self.pat = regex.compile( - r"<\|startoftext\|>|<\|endoftext\|>|\'s|\'t|\'re|\'ve|\'m|\'ll|\'d|" - r"[\p{L}]+|[\p{N}]|[^\s\p{L}\p{N}]+", - regex.IGNORECASE, - ) - - def _basic_clean(self, text): - return html.unescape(html.unescape(self.ftfy.fix_text(text))).strip() - - def _whitespace_clean(self, text): - return self.regex.sub(r"\s+", " ", text).strip() - - def bpe(self, token): - if token in self.cache: - return self.cache[token] - word = tuple(token[:-1]) + (token[-1] + "",) - pairs = _get_pairs(word) - if not pairs: - return token + "" - while True: - bigram = min(pairs, key=lambda p: self.bpe_ranks.get(p, float("inf"))) - if bigram not in self.bpe_ranks: - break - first, second = bigram - new_word = [] - i = 0 - while i < len(word): - try: - j = word.index(first, i) - except ValueError: - new_word.extend(word[i:]) - break - new_word.extend(word[i:j]) - if word[j] == first and j < len(word) - 1 and word[j + 1] == second: - new_word.append(first + second) - i = j + 2 - else: - new_word.append(word[j]) - i = j + 1 - word = tuple(new_word) - if len(word) == 1: - break - pairs = _get_pairs(word) - out = " ".join(word) - self.cache[token] = out - return out - - def encode(self, text): - text = self._whitespace_clean(self._basic_clean(text)).lower() - tokens = [] - for token in self.regex.findall(self.pat, text): - token = "".join(self.byte_encoder[b] for b in token.encode("utf-8")) - tokens.extend(self.encoder[piece] for piece in self.bpe(token).split(" ")) - return tokens - - -class CLIPQuickGELU(nn.Module): - def forward(self, x): - return x * torch.sigmoid(1.702 * x) - - -class CLIPTextBlock(nn.Module): - def __init__(self, width, heads, mask): - super().__init__() - self.attn = nn.MultiheadAttention(width, heads) - self.ln_1 = nn.LayerNorm(width) - self.mlp = nn.Sequential() - self.mlp.add_module("c_fc", nn.Linear(width, width * 4)) - self.mlp.add_module("gelu", CLIPQuickGELU()) - self.mlp.add_module("c_proj", nn.Linear(width * 4, width)) - self.ln_2 = nn.LayerNorm(width) - self.register_buffer("mask", mask, persistent=False) - - def forward(self, x): - x_ln = self.ln_1(x) - x = x + self.attn(x_ln, x_ln, x_ln, need_weights=False, attn_mask=self.mask)[0] - x = x + self.mlp(self.ln_2(x)) - return x - - -class CLIPTextTower(nn.Module): - def __init__(self, state_dict): - super().__init__() - embed_dim = state_dict["text_projection"].shape[1] - context_length = state_dict["positional_embedding"].shape[0] - vocab_size = state_dict["token_embedding.weight"].shape[0] - width = state_dict["ln_final.weight"].shape[0] - heads = width // 64 - layers = len( - {k.split(".")[2] for k in state_dict if k.startswith("transformer.resblocks")} - ) - mask = torch.empty(context_length, context_length).fill_(float("-inf")).triu_(1) - - self.context_length = context_length - self.token_embedding = nn.Embedding(vocab_size, width) - self.positional_embedding = nn.Parameter(torch.empty(context_length, width)) - self.transformer = nn.Module() - self.transformer.resblocks = nn.Sequential( - *[CLIPTextBlock(width, heads, mask) for _ in range(layers)] - ) - self.ln_final = nn.LayerNorm(width) - self.text_projection = nn.Parameter(torch.empty(width, embed_dim)) - - def forward(self, text): - x = self.token_embedding(text) + self.positional_embedding - x = x.permute(1, 0, 2) - x = self.transformer.resblocks(x) - x = x.permute(1, 0, 2) - x = self.ln_final(x) - return x[torch.arange(x.shape[0], device=x.device), text.argmax(dim=-1)] @ self.text_projection - - -class CLIPTextEmbedder: - """OpenAI CLIP ViT-B/32 text encoder that returns normalized 512-d vectors. - Holds a CPU memo cache keyed by prompt so repeats skip BPE+transformer.""" - - def __init__(self, config: ClipConfig, device="cpu"): - b32_path, bpe_path = ensure_clip_text_assets(config) - self.device = torch.device(device) - try: - state_dict = torch.jit.load(str(b32_path), map_location="cpu").state_dict() - except RuntimeError: - state_dict = torch.load(b32_path, map_location="cpu", weights_only=False) - self.tokenizer = CLIPBPETokenizer(bpe_path) - self.model = CLIPTextTower(state_dict).eval().to(self.device) - text_keys = { - k: v - for k, v in state_dict.items() - if k.startswith( - ( - "token_embedding", - "positional_embedding", - "transformer.resblocks", - "ln_final", - "text_projection", - ) - ) - } - missing, unexpected = self.model.load_state_dict(text_keys, strict=False) - if missing or unexpected: - raise RuntimeError(f"bad CLIP text weights: missing={missing} unexpected={unexpected}") - self._cache = {} - - @torch.no_grad() - def encode(self, texts): - if isinstance(texts, str): - texts = [texts] - fresh = [t for t in dict.fromkeys(texts) if t not in self._cache] - if fresh: - context = torch.zeros( - len(fresh), self.model.context_length, dtype=torch.long, device=self.device - ) - for i, text in enumerate(fresh): - token_ids = [ - self.tokenizer.encoder[SOT_TOKEN], - *self.tokenizer.encode(text), - self.tokenizer.encoder[EOT_TOKEN], - ] - if len(token_ids) > self.model.context_length: - raise RuntimeError( - f"Input {text!r} is too long for CLIP context length " - f"{self.model.context_length}" - ) - context[i, : len(token_ids)] = torch.tensor( - token_ids, dtype=torch.long, device=self.device - ) - features = self.model(context) - features = features / features.norm(dim=-1, keepdim=True) - for i, text in enumerate(fresh): - self._cache[text] = features[i].cpu() - out = torch.stack([self._cache[t] for t in texts], dim=0) - return out.to(self.device) - - -def encode_clip_text(texts, config: ClipConfig, device="cpu"): - """Encode exact prompt text with OpenAI CLIP ViT-B/32.""" - return CLIPTextEmbedder(config, device=device).encode(texts) - - -def encode_clip_task_name(task_names, config: ClipConfig, device="cpu"): - """Encode task names after production-style dash/underscore replacement.""" - if isinstance(task_names, str): - task_names = [task_names] - return encode_clip_text([task_name_to_prompt(t) for t in task_names], config, device=device) - - -# DINOv3 ViT-B/16 vision encoder. - - -def _rope_rotate_half(x): - x1, x2 = x.chunk(2, dim=-1) - return torch.cat([-x2, x1], dim=-1) - - -class DinoRope(nn.Module): - """RoPE over the 2D patch grid (base=100, separate coord normalization). - - rescale_coords=2 applies a random log-uniform rescale of the coordinates - during training only — part of the pretraining distribution, kept for - finetuning fidelity. - """ - - def __init__(self, embed_dim, num_heads, base=100.0, rescale_coords=2.0): - super().__init__() - d_head = embed_dim // num_heads - self.d_head = d_head - self.rescale_coords = rescale_coords - self.register_buffer("periods", torch.empty(d_head // 4), persistent=True) - with torch.no_grad(): - self.periods.copy_( - base ** (2 * torch.arange(d_head // 4, dtype=torch.float32) / (d_head // 2)) - ) - - def forward(self, H, W): - dev = self.periods.device - coords_h = torch.arange(0.5, H, device=dev, dtype=torch.float32) / H - coords_w = torch.arange(0.5, W, device=dev, dtype=torch.float32) / W - coords = torch.stack(torch.meshgrid(coords_h, coords_w, indexing="ij"), dim=-1) - coords = coords.flatten(0, 1) - coords = 2.0 * coords - 1.0 - if self.training and self.rescale_coords is not None: - r = np.log(self.rescale_coords) - rescale = torch.empty(1, device=dev).uniform_(-r, r).exp() - coords = coords * rescale - angles = 2 * math.pi * coords[:, :, None] / self.periods[None, None, :] - angles = angles.flatten(1, 2).tile(2) - return torch.sin(angles), torch.cos(angles) - - -class LinearKMaskedBias(nn.Linear): - """qkv Linear whose k-third of the bias is masked to zero (DINOv3 quirk).""" - - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - self.register_buffer("bias_mask", torch.full_like(self.bias, math.nan)) - - def forward(self, x): - return F.linear(x, self.weight, self.bias * self.bias_mask.to(self.bias.dtype)) - - -class DinoAttention(nn.Module): - def __init__(self, dim, num_heads): - super().__init__() - self.num_heads = num_heads - self.qkv = LinearKMaskedBias(dim, dim * 3, bias=True) - self.proj = nn.Linear(dim, dim, bias=True) - - def forward(self, x, rope=None): - B, N, C = x.shape - qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads) - q, k, v = [t.transpose(1, 2) for t in torch.unbind(qkv, 2)] - if rope is not None: - sin, cos = rope - n_prefix = N - sin.shape[-2] # cls + storage tokens are not rotated - q_dt, k_dt = q.dtype, k.dtype - q, k = q.to(sin.dtype), k.to(sin.dtype) - q = torch.cat( - [q[:, :, :n_prefix], q[:, :, n_prefix:] * cos + _rope_rotate_half(q[:, :, n_prefix:]) * sin], - dim=-2, - ) - k = torch.cat( - [k[:, :, :n_prefix], k[:, :, n_prefix:] * cos + _rope_rotate_half(k[:, :, n_prefix:]) * sin], - dim=-2, - ) - q, k = q.to(q_dt), k.to(k_dt) - x = F.scaled_dot_product_attention(q, k, v) - return self.proj(x.transpose(1, 2).reshape(B, N, C)) - - -class LayerScale(nn.Module): - def __init__(self, dim, init_values=1e-5): - super().__init__() - self.gamma = nn.Parameter(init_values * torch.ones(dim)) - - def forward(self, x): - return x * self.gamma - - -class DinoMlp(nn.Module): - def __init__(self, dim, hidden): - super().__init__() - self.fc1 = nn.Linear(dim, hidden) - self.act = nn.GELU() - self.fc2 = nn.Linear(hidden, dim) - - def forward(self, x): - return self.fc2(self.act(self.fc1(x))) - - -class DinoBlock(nn.Module): - def __init__(self, dim, num_heads, ffn_ratio=4.0): - super().__init__() - self.norm1 = nn.LayerNorm(dim, eps=1e-5) - self.attn = DinoAttention(dim, num_heads) - self.ls1 = LayerScale(dim) - self.norm2 = nn.LayerNorm(dim, eps=1e-5) - self.mlp = DinoMlp(dim, int(dim * ffn_ratio)) - self.ls2 = LayerScale(dim) - - def forward(self, x, rope=None): - x = x + self.ls1(self.attn(self.norm1(x), rope=rope)) - x = x + self.ls2(self.mlp(self.norm2(x))) - return x - - -class DinoPatchEmbed(nn.Module): - def __init__(self, patch_size=16, in_chans=3, embed_dim=768): - super().__init__() - self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size) - - def forward(self, x): - x = self.proj(x) # (B, D, H/16, W/16) - return x.flatten(2).transpose(1, 2), x.shape[2], x.shape[3] - - -class DinoVisionTransformer(nn.Module): - """DINOv3 ViT-B/16 with 4 storage tokens. encode_image_tokens() returns - (B, 1+196, 768) = CLS + patch tokens (storage tokens dropped), matching - the production vision backbone interface.""" - - N_STORAGE_TOKENS = 4 - - def __init__(self, embed_dim, depth, num_heads): - super().__init__() - self.embed_dim = embed_dim - self.patch_embed = DinoPatchEmbed(embed_dim=embed_dim) - self.cls_token = nn.Parameter(torch.empty(1, 1, embed_dim)) - self.storage_tokens = nn.Parameter(torch.empty(1, self.N_STORAGE_TOKENS, embed_dim)) - self.mask_token = nn.Parameter(torch.empty(1, embed_dim)) - self.rope_embed = DinoRope(embed_dim, num_heads) - self.blocks = nn.ModuleList(DinoBlock(embed_dim, num_heads) for _ in range(depth)) - self.norm = nn.LayerNorm(embed_dim, eps=1e-5) - self.init_weights() - - def init_weights(self): - """Match production's models/dinov3/vision_transformer.py:init_weights_vit. - Crucially, this fills `bias_mask` (otherwise NaN-initialized) so that - the K-third of every qkv bias is masked to 0 — without this, a fresh - DINOv3 produces NaN on its very first forward pass.""" - nn.init.normal_(self.cls_token, std=0.02) - nn.init.normal_(self.storage_tokens, std=0.02) - nn.init.zeros_(self.mask_token) - for m in self.modules(): - if isinstance(m, nn.Linear): - nn.init.trunc_normal_(m.weight, std=0.02) - if m.bias is not None: - nn.init.zeros_(m.bias) - if isinstance(m, LinearKMaskedBias): - o = m.out_features - m.bias_mask.fill_(1) - m.bias_mask[o // 3 : 2 * o // 3].fill_(0) - elif isinstance(m, nn.LayerNorm): - m.reset_parameters() - elif isinstance(m, LayerScale): - nn.init.constant_(m.gamma, 1e-5) - elif isinstance(m, DinoPatchEmbed): - # Match nn.Conv2d default - m.proj.reset_parameters() - - def encode_image_tokens(self, images): - x, H, W = self.patch_embed(images) - B = x.shape[0] - cls_token = self.cls_token + 0 * self.mask_token # production quirk, kept - x = torch.cat( - [cls_token.expand(B, -1, -1), self.storage_tokens.expand(B, -1, -1), x], dim=1 - ) - rope = self.rope_embed(H, W) - for blk in self.blocks: - x = blk(x, rope=rope) - x = self.norm(x) - cls_out = x[:, :1] - patches = x[:, 1 + self.N_STORAGE_TOKENS :] - return torch.cat([cls_out, patches], dim=1) - - -class DinoVisionBackbone(nn.Module): - """Wrapper around DinoVisionTransformer with an optional bf16-autocast - forward path. The wrapper keeps the production checkpoint key layout - (`img_backbone.dinov3_model.*`) so the slim 200k checkpoint loads with - zero missing/unexpected keys. Set bf16 with set_bfloat16(True): the - DINO forward then runs under autocast(bf16) on CUDA, cutting - vision-encoder activation memory roughly in half. Tokens are cast back - to fp32 on the way out so the surrounding DiT stays dtype-stable. - """ - - def __init__(self, config: DiTConfig): - super().__init__() - self.dinov3_model = DinoVisionTransformer( - embed_dim=config.vit_embed_dim, - depth=config.vit_depth, - num_heads=config.vit_num_heads, - ) - self.bfloat16 = False - - def set_bfloat16(self, enabled: bool = True): - self.bfloat16 = bool(enabled) - - def encode_image_tokens(self, images): - if self.bfloat16 and images.is_cuda: - with torch.autocast(device_type="cuda", dtype=torch.bfloat16): - tokens = self.dinov3_model.encode_image_tokens(images) - return tokens.to(torch.float32) - return self.dinov3_model.encode_image_tokens(images) - - -# ABC-DiT policy. - - -def modulate(x, shift, scale): - if shift.ndim == 2: - shift = shift.unsqueeze(1) - scale = scale.unsqueeze(1) - return x * (1 + scale) + shift - - -def gate_residual(gate, residual): - if gate.ndim == 2: - gate = gate.unsqueeze(1) - return gate * residual - - -def get_1d_sincos_pos_embed(embed_dim, length): - omega = np.arange(embed_dim // 2, dtype=np.float64) - omega /= embed_dim / 2.0 - omega = 1.0 / 10000**omega - out = np.einsum("m,d->md", np.arange(length, dtype=np.float64), omega) - return np.concatenate([np.sin(out), np.cos(out)], axis=1) - - -class TimestepEmbedder(nn.Module): - def __init__(self, hidden_size, frequency_embedding_size=256): - super().__init__() - self.mlp = nn.Sequential( - nn.Linear(frequency_embedding_size, hidden_size, bias=True), - nn.SiLU(), - nn.Linear(hidden_size, hidden_size, bias=True), - ) - self.frequency_embedding_size = frequency_embedding_size - half = frequency_embedding_size // 2 - freqs = torch.exp( - -math.log(10000) * torch.arange(half, dtype=torch.float32) / half - ) - self.register_buffer("freqs", freqs, persistent=False) - - def timestep_embedding(self, t): - freqs = self.freqs - if freqs.device != t.device: - freqs = freqs.to(device=t.device) - args = t[:, None].float() * freqs[None] - return torch.cat([torch.cos(args), torch.sin(args)], dim=-1) - - def forward(self, t): - t_shape = t.shape - t_freq = self.timestep_embedding(t.reshape(-1)) - t_emb = self.mlp(t_freq.to(self.mlp[0].weight.dtype)) - return t_emb.reshape(*t_shape, -1) - - -class DiTAttention(nn.Module): - """Self-attention over action tokens (timm-equivalent, qkv_bias=True).""" - - def __init__(self, dim, num_heads): - super().__init__() - self.num_heads = num_heads - self.qkv = nn.Linear(dim, dim * 3, bias=True) - self.proj = nn.Linear(dim, dim, bias=True) - - def forward(self, x): - B, N, C = x.shape - qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads) - q, k, v = qkv.permute(2, 0, 3, 1, 4).unbind(0) - x = F.scaled_dot_product_attention(q, k, v) - return self.proj(x.transpose(1, 2).reshape(B, N, C)) - - -class DiTMlp(nn.Module): - def __init__(self, dim, hidden): - super().__init__() - self.fc1 = nn.Linear(dim, hidden) - self.act = nn.GELU(approximate="tanh") - self.fc2 = nn.Linear(hidden, dim) - - def forward(self, x): - return self.fc2(self.act(self.fc1(x))) - - -class DiTBlock(nn.Module): - """AdaLN-Zero DiT block with vision cross-attention (9-way modulation).""" - - def __init__(self, hidden_size, num_heads, mlp_ratio=4.0): - super().__init__() - self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) - self.attn = DiTAttention(hidden_size, num_heads) - self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) - self.mlp = DiTMlp(hidden_size, int(hidden_size * mlp_ratio)) - self.norm_xattn = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) - self.norm_xattn_kv = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) - self.cross_attn = nn.MultiheadAttention(hidden_size, num_heads, batch_first=True) - self.adaLN_modulation = nn.Sequential( - nn.SiLU(), nn.Linear(hidden_size, 9 * hidden_size, bias=True) - ) - - def forward(self, x, c, vision_tokens): - ( - shift_msa, scale_msa, gate_msa, - shift_xattn, scale_xattn, gate_xattn, - shift_mlp, scale_mlp, gate_mlp, - ) = self.adaLN_modulation(c).chunk(9, dim=-1) - - x = x + gate_residual(gate_msa, self.attn(modulate(self.norm1(x), shift_msa, scale_msa))) - - x_normed = modulate(self.norm_xattn(x), shift_xattn, scale_xattn) - kv = self.norm_xattn_kv(vision_tokens) - xattn_out, _ = self.cross_attn(x_normed, kv, kv, need_weights=False) - x = x + gate_residual(gate_xattn, xattn_out) - - x = x + gate_residual(gate_mlp, self.mlp(modulate(self.norm2(x), shift_mlp, scale_mlp))) - return x - - -class FinalLayer(nn.Module): - def __init__(self, hidden_size, action_dim): - super().__init__() - self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) - self.linear = nn.Linear(hidden_size, action_dim, bias=True) - self.adaLN_modulation = nn.Sequential( - nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size, bias=True) - ) - - def forward(self, x, c): - shift, scale = self.adaLN_modulation(c).chunk(2, dim=-1) - return self.linear(modulate(self.norm_final(x), shift, scale)) - - -class PoolMlp(nn.Module): - def __init__(self, in_dim, hidden_dim): - super().__init__() - self.fc1 = nn.Linear(in_dim, hidden_dim) - self.act = nn.GELU() - self.fc2 = nn.Linear(hidden_dim, in_dim) - - def forward(self, x): - return self.fc2(self.act(self.fc1(x))) - - -class AttentionPoolBlock(nn.Module): - """Learnable queries cross-attend to ViT tokens (per camera).""" - - def __init__(self, embed_dim, num_heads, mlp_ratio=4): - super().__init__() - self.ln_1 = nn.LayerNorm(embed_dim) - self.attention = nn.MultiheadAttention(embed_dim, num_heads, batch_first=True) - self.ln_2 = nn.LayerNorm(embed_dim) - self.mlp = PoolMlp(embed_dim, int(mlp_ratio * embed_dim)) - - def forward(self, x, queries): - x_kv = self.ln_1(x) - x_q = self.ln_1(queries) - out, _ = self.attention(x_q, x_kv, x_kv, need_weights=False) - return self.mlp(self.ln_2(out)) + out - - -class DiTPolicy(nn.Module): - """Minimal ABC-DiT: loads the production dit_xL pretraining checkpoint.""" - - def __init__(self, config: DiTConfig): - super().__init__() - self.config = config - H = config.hidden_size - self.camera_keys = list(config.camera_keys) - self.chunk_length = config.chunk_length - self.action_dim = config.action_dim - - self.x_embedder = nn.Linear(config.state_dim, H) - self.y_embedder = nn.Linear(config.action_dim, H) - # Checkpoint compatibility only; unused in forward. - self.img_proj = nn.Linear(config.vit_embed_dim, H) - self.img_proj.requires_grad_(False) - self.t_embedder = TimestepEmbedder(H) - self.pos_embed = nn.Parameter(torch.zeros(1, config.chunk_length, H), requires_grad=False) - - self.img_backbone = DinoVisionBackbone(config) - - self.apool_queries = nn.ParameterDict( - { - cam: nn.Parameter( - torch.randn(1, config.vision_pool_num_queries, config.vit_embed_dim) * 0.02 - ) - for cam in self.camera_keys - } - ) - self.apool = nn.ModuleDict( - { - cam: AttentionPoolBlock( - config.vit_embed_dim, - config.vision_pool_num_heads, - config.vision_pool_mlp_ratio, - ) - for cam in self.camera_keys - } - ) - self.vision_tokens_proj = nn.Linear(config.vit_embed_dim, H) - self.vision_camera_embed = nn.Embedding(len(self.camera_keys), H) - - self.task_to_hidden = nn.Linear(config.task_embed_dim, H) - self.blocks = nn.ModuleList( - DiTBlock(H, config.num_heads, config.mlp_ratio) for _ in range(config.depth) - ) - self.final_layer = FinalLayer(H, config.action_dim) - - # cond = [state, task, timestep] -> hidden (vision goes via cross-attn) - self.cond_proj = nn.Sequential( - nn.Linear(3 * H, H), nn.SiLU(), nn.Linear(H, H), nn.LayerNorm(H) - ) - - self.register_buffer("clip_mean", torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1)) - self.register_buffer("clip_std", torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1)) - - pos = get_1d_sincos_pos_embed(H, config.chunk_length) - self.pos_embed.data.copy_(torch.from_numpy(pos).float().unsqueeze(0)) - - def build_vision_tokens(self, images): - """images: dict cam -> (B, 3, 224, 224), already ImageNet-normalized. - Returns (B, num_cameras * queries, hidden).""" - pooled = [] - for cam in self.camera_keys: - tokens = self.img_backbone.encode_image_tokens(images[cam]) - tokens = tokens.to(self.apool_queries[cam].dtype) - queries = self.apool_queries[cam].expand(tokens.shape[0], -1, -1) - pooled.append(self.apool[cam](tokens, queries)) - tokens_by_camera = torch.stack(pooled, dim=1) # (B, Nc, K, vit_dim) - B, Nc, K, D = tokens_by_camera.shape - vision_tokens = self.vision_tokens_proj( - tokens_by_camera.reshape(B * Nc * K, D).to(self.vision_tokens_proj.weight.dtype) - ).reshape(B, Nc, K, -1) - cam_emb = self.vision_camera_embed(torch.arange(Nc, device=vision_tokens.device)) - vision_tokens = vision_tokens + cam_emb[None, :, None, :] - return vision_tokens.reshape(B, Nc * K, -1) - - def compute_cond(self, state, task_vec_clip, t_cond): - """state (B,14); task_vec_clip (B,512); t_cond (B,) or (B,T). - Returns conditioning c: (B,H) or (B,T,H).""" - model_dtype = self.x_embedder.weight.dtype - cond_dtype = self.cond_proj[0].weight.dtype - st_vec = self.x_embedder(state.to(model_dtype)) - task_vec_h = self.task_to_hidden(task_vec_clip.to(self.task_to_hidden.weight.dtype)) - task_vec_h = task_vec_h.to(model_dtype) - t_vec = self.t_embedder(t_cond.to(model_dtype)) - cond_parts = [st_vec, task_vec_h, t_vec] - if t_vec.ndim == 3: - T = t_vec.shape[1] - cond_parts = [ - p.unsqueeze(1).expand(-1, T, -1) if p.ndim == 2 else p for p in cond_parts - ] - cond_concat = torch.cat(cond_parts, dim=-1).to(cond_dtype) - if cond_dtype == torch.float32 and cond_concat.is_cuda: - with torch.autocast(device_type="cuda", enabled=False): - return self.cond_proj(cond_concat).to(model_dtype) - return self.cond_proj(cond_concat).to(model_dtype) - - def predict_velocity(self, x_t, c, vision_tokens): - z = self.y_embedder(x_t) + self.pos_embed.data[:, : x_t.shape[1], :] - for block in self.blocks: - z = block(z, c, vision_tokens) - return self.final_layer(z, c) - - def forward( - self, - batch, - noise=None, - t=None, - max_action_prefix=0, - prefix_conditioning_prob=1.0, - prefix_noise_scale=0.0, - ): - """Flow-matching training loss with optional action-prefix conditioning. - batch: state (B,14), images dict, actions (B,30,14), task_vec_clip (B,512), - optional state_is_masked (B,) bool.""" - state = batch["state"] - actions = batch["actions"] - N, T_chunk, D_action = actions.shape - - if noise is None: - noise = torch.randn_like(actions) - if t is None: - t = torch.rand(N, 1, 1, device=state.device, dtype=actions.dtype) - - if max_action_prefix > 0: - apply_prefix = torch.rand(N, device=state.device) < prefix_conditioning_prob - if "state_is_masked" in batch: - apply_prefix = apply_prefix & ~batch["state_is_masked"].to(state.device) - delay = torch.randint(0, max_action_prefix, (N,), device=state.device) - delay = torch.where(apply_prefix, delay, torch.zeros_like(delay)) - prefix_mask = torch.arange(T_chunk, device=state.device)[None, :] < delay[:, None] - prefix_mask_expanded = prefix_mask.unsqueeze(-1) - t_per_pos = torch.where(prefix_mask_expanded, torch.zeros_like(t), t) - else: - prefix_mask_expanded = None - t_per_pos = t - - x_t = (1 - t_per_pos) * actions + t_per_pos * noise - if prefix_noise_scale > 0.0 and prefix_mask_expanded is not None: - x_t = x_t + prefix_mask_expanded.float() * torch.randn_like(x_t) * prefix_noise_scale - - vision_tokens = self.build_vision_tokens(batch["images"]) - t_cond = t_per_pos.squeeze(-1) if prefix_mask_expanded is not None else t[:, 0, 0] - c = self.compute_cond(state, batch["task_vec_clip"], t_cond) - v_t = self.predict_velocity(x_t, c, vision_tokens) - - u_t = noise - actions - if prefix_mask_expanded is not None: - postfix_mask = ~prefix_mask_expanded - masked_loss = ((u_t - v_t) ** 2) * postfix_mask.float() - return masked_loss.sum() / (postfix_mask.float().sum() * D_action + 1e-8) - return F.mse_loss(u_t, v_t) - - @torch.no_grad() - def sample_actions(self, batch, num_steps=10, noise=None): - """Euler flow integration from noise to actions (production tau=1 path). - Vision tokens and the static conditioning are computed once and reused - across steps, like production infer().""" - state = batch["state"] - B = state.shape[0] - model_dtype = self.y_embedder.weight.dtype - if noise is None: - noise = torch.randn( - B, - self.chunk_length, - self.action_dim, - device=state.device, - dtype=model_dtype, - ) - x_t = noise.to(device=state.device, dtype=model_dtype) - vision_tokens = self.build_vision_tokens(batch["images"]) - dt = -1.0 / num_steps - for i in range(num_steps): - t = torch.full((B,), 1.0 + i * dt, device=state.device, dtype=model_dtype) - c = self.compute_cond(state, batch["task_vec_clip"], t) - v = self.predict_velocity(x_t, c, vision_tokens) - x_t = x_t + v * dt - return x_t - - @torch.no_grad() - def sample_actions_rtc(self, batch, action_prefix, prefix_length: int, num_steps=10, noise=None): - """Euler sampling with per-position action-prefix conditioning.""" - state = batch["state"] - B = state.shape[0] - model_dtype = self.y_embedder.weight.dtype - if noise is None: - noise = torch.randn( - B, - self.chunk_length, - self.action_dim, - device=state.device, - dtype=model_dtype, - ) - x_t = noise.to(device=state.device, dtype=model_dtype) - action_prefix = action_prefix.to(device=state.device, dtype=model_dtype) - - prefix_pos = torch.arange(self.chunk_length, device=state.device) < prefix_length - prefix_mask = prefix_pos.view(1, self.chunk_length, 1).expand_as(x_t) - prefix_t_mask = prefix_pos.view(1, self.chunk_length).expand(B, self.chunk_length) - x_t = torch.where(prefix_mask, action_prefix, x_t) - - vision_tokens = self.build_vision_tokens(batch["images"]) - dt = -1.0 / num_steps - for i in range(num_steps): - t = torch.full( - (B, self.chunk_length), - 1.0 + i * dt, - device=state.device, - dtype=model_dtype, - ) - t = torch.where(prefix_t_mask, torch.zeros_like(t), t) - c = self.compute_cond(state, batch["task_vec_clip"], t) - v = self.predict_velocity(x_t, c, vision_tokens) - x_t = x_t + v * dt - x_t = torch.where(prefix_mask, action_prefix, x_t) - return x_t - - -def load_pretrained(model, ckpt_path): - """Load the slim production checkpoint (model-only, prefixes stripped).""" - ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False, mmap=True) - sd = ckpt["model"] if "model" in ckpt else ckpt - sd = {k[len("_orig_mod."):] if k.startswith("_orig_mod.") else k: v for k, v in sd.items()} - missing, unexpected = model.load_state_dict(sd, strict=False) - if unexpected: - raise RuntimeError(f"unexpected checkpoint keys: {unexpected[:8]}") - if missing: - raise RuntimeError(f"missing checkpoint keys: {missing[:8]}") - return ckpt - - -if __name__ == "__main__": - model = DiTPolicy(DiTConfig()) - n_params = sum(p.numel() for p in model.parameters()) - print(f"DiTPolicy built: {n_params / 1e9:.3f}B params") diff --git a/dimos/imitation/policy/abc/python/abc_minimal/fast_inference.py b/dimos/imitation/policy/abc/python/abc_minimal/fast_inference.py deleted file mode 100644 index 42031c8cc3..0000000000 --- a/dimos/imitation/policy/abc/python/abc_minimal/fast_inference.py +++ /dev/null @@ -1,118 +0,0 @@ -"""CUDA graph helpers for fast ABC-DiT policy inference.""" - -from __future__ import annotations - -from typing import Any, Protocol - -import numpy as np -import torch - -from abc_minimal.preprocess import normalize, resize_pad_normalize, unnormalize - - -class _PolicyForFastInference(Protocol): - model: Any - device: torch.device - config: Any - task_vec: torch.Tensor - norm_stats: dict[str, Any] - diffusion_steps: int - -class FastInferenceGraph: - """Fixed-shape outer CUDA graph over `model.sample_actions`. - - Captures one bf16 batch on GPU; subsequent .infer() calls just memcpy - fresh inputs into the static tensors and replay the graph. - """ - - def __init__(self, policy: _PolicyForFastInference): - self.policy = policy - self.model = policy.model - self.device = policy.device - self.dtype = torch.bfloat16 - self.graph = torch.cuda.CUDAGraph() - self.output: torch.Tensor | None = None - m = policy.config.model - - self.static_state = torch.empty(1, m.state_dim, device=self.device, dtype=self.dtype) - self.static_noise = torch.empty( - 1, m.chunk_length, m.action_dim, device=self.device, dtype=self.dtype - ) - self.static_images = { - cam: torch.empty(1, 3, 224, 224, device=self.device, dtype=self.dtype) - for cam in m.camera_keys - } - self.static_task_vec = policy.task_vec.to(device=self.device, dtype=self.dtype).clone() - self.batch = { - "state": self.static_state, - "actions": torch.zeros( - 1, m.chunk_length, m.action_dim, device=self.device, dtype=self.dtype - ), - "images": self.static_images, - "task_vec_clip": self.static_task_vec, - } - - def _copy_inputs(self, obs: dict[str, Any], noise: np.ndarray | None) -> None: - m = self.policy.config.model - state = normalize( - np.asarray(obs["state"], dtype=np.float32), self.policy.norm_stats["state"] - ) - self.static_state.copy_( - torch.from_numpy(state[None]).to(device=self.device, dtype=self.dtype) - ) - if noise is None: - self.static_noise.normal_() - else: - noise_arr = noise[None].astype(np.float32, copy=False) - if noise_arr.shape != (1, m.chunk_length, m.action_dim): - raise ValueError( - f"fast inference expects noise shape " - f"{(m.chunk_length, m.action_dim)}, got {noise.shape}" - ) - self.static_noise.copy_( - torch.from_numpy(noise_arr).to(device=self.device, dtype=self.dtype) - ) - for cam in m.camera_keys: - self.static_images[cam].copy_( - resize_pad_normalize(obs["images"][cam]) - .unsqueeze(0) - .to(device=self.device, dtype=self.dtype) - ) - - def capture( - self, - warmup_obs: dict[str, Any], - warmup_noise: np.ndarray | None, - replay_warmups: int, - ) -> None: - self._copy_inputs(warmup_obs, warmup_noise) - - stream = torch.cuda.Stream() - stream.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(stream): - for _ in range(5): - self.output = self.model.sample_actions( - self.batch, - num_steps=self.policy.diffusion_steps, - noise=self.static_noise, - ) - torch.cuda.current_stream().wait_stream(stream) - - with torch.cuda.graph(self.graph): - self.output = self.model.sample_actions( - self.batch, num_steps=self.policy.diffusion_steps, noise=self.static_noise - ) - - for _ in range(replay_warmups): - self._copy_inputs(warmup_obs, warmup_noise) - self.graph.replay() - assert self.output is not None - _ = self.output[0].float().detach().cpu().numpy() - torch.cuda.synchronize() - - def infer(self, obs: dict[str, Any], noise: np.ndarray | None) -> np.ndarray: - self._copy_inputs(obs, noise) - self.graph.replay() - assert self.output is not None - actions_np = self.output[0].float().detach().cpu().numpy() - return unnormalize(actions_np, self.policy.norm_stats["actions"]).astype(np.float32) diff --git a/dimos/imitation/policy/abc/python/abc_minimal/preprocess.py b/dimos/imitation/policy/abc/python/abc_minimal/preprocess.py deleted file mode 100644 index a1306afc37..0000000000 --- a/dimos/imitation/policy/abc/python/abc_minimal/preprocess.py +++ /dev/null @@ -1,71 +0,0 @@ -"""Shared state/action normalization and image preprocessing.""" - -import json -from pathlib import Path - -import numpy as np -import torch -import torch.nn.functional as F - - -IMAGENET_MEAN = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1) -IMAGENET_STD = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1) - - -def parse_norm_stats(raw): - stats = raw.get("norm_stats", raw) - if "state" not in stats and "actions" not in stats: - key = "xdof" if "xdof" in stats else next(iter(stats)) - stats = stats[key] - return { - key: {k: np.asarray(v, dtype=np.float32) for k, v in stats[key].items()} - for key in ("state", "actions") - } - - -def load_norm_stats(path): - return parse_norm_stats(json.loads(Path(path).read_text())) - - -def normalize(x, stats): - return (x - stats["mean"]) / (stats["std"] + 1e-6) - - -def unnormalize(x, stats): - return x * (stats["std"] + 1e-6) + stats["mean"] - - -def resize_with_pad(img_hwc, target_h=224, target_w=224): - h, w, _ = img_hwc.shape - if (h, w) == (target_h, target_w): - return img_hwc - ratio = max(w / target_w, h / target_h) - new_h = max(1, int(round(h / ratio))) - new_w = max(1, int(round(w / ratio))) - resized = F.interpolate( - img_hwc.permute(2, 0, 1).unsqueeze(0), - size=(new_h, new_w), - mode="bilinear", - align_corners=False, - antialias=True, - ).squeeze(0) - pad_h0 = (target_h - new_h) // 2 - pad_h1 = target_h - new_h - pad_h0 - pad_w0 = (target_w - new_w) // 2 - pad_w1 = target_w - new_w - pad_w0 - padded = F.pad(resized, (pad_w0, pad_w1, pad_h0, pad_h1), value=0) - return padded.permute(1, 2, 0) - - -def imagenet_normalize(img_chw): - mean = IMAGENET_MEAN.to(device=img_chw.device, dtype=img_chw.dtype) - std = IMAGENET_STD.to(device=img_chw.device, dtype=img_chw.dtype) - return (img_chw - mean) / (std + 1e-6) - - -def resize_pad_normalize(img_chw, target_h=224, target_w=224): - x = torch.as_tensor(img_chw).float() - if x.max() > 1.0: - x = x / 255.0 - x = resize_with_pad(x.permute(1, 2, 0), target_h, target_w).permute(2, 0, 1) - return imagenet_normalize(x) diff --git a/dimos/imitation/policy/abc/python/dimos_abc/__init__.py b/dimos/imitation/policy/abc/python/dimos_abc/__init__.py deleted file mode 100644 index e908044542..0000000000 --- a/dimos/imitation/policy/abc/python/dimos_abc/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""DimOS adapter for the vendored ABC-DiT inference code.""" diff --git a/dimos/imitation/policy/abc/python/dimos_abc/runtime.py b/dimos/imitation/policy/abc/python/dimos_abc/runtime.py deleted file mode 100644 index 1c8d21e2d6..0000000000 --- a/dimos/imitation/policy/abc/python/dimos_abc/runtime.py +++ /dev/null @@ -1,208 +0,0 @@ -# Copyright 2026 Dimensional Inc. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Amazon ABC-DiT adapter for the shared DimOS rollout runtime.""" - -from __future__ import annotations - -from collections.abc import Mapping -import json -from pathlib import Path -from types import SimpleNamespace -from typing import Any, cast - -from abc_minimal.config import ClipConfig, DiTConfig -from abc_minimal.dit import CLIPTextEmbedder, DiTPolicy, load_pretrained -from abc_minimal.fast_inference import FastInferenceGraph -from abc_minimal.preprocess import normalize, parse_norm_stats, resize_pad_normalize, unnormalize -import numpy as np -from numpy.typing import NDArray -import torch - -from dimos.imitation.policy.abc.module import AbcPolicyConfig, DualOpenYamAbcPolicy -from dimos.imitation.policy.backend import PolicyBackendInfo -from dimos.imitation.policy.runtime import declare_policy_runtime -from dimos.imitation.profile import ImageSource, JointPositionSource, PolicyIOProfile - -torch.set_float32_matmul_precision("high") - - -class AbcBackend: - """Run the released 14-D, three-camera ABC-DiT policy in process.""" - - def __init__(self, config: AbcPolicyConfig) -> None: - self._rollout_config = config - self.config = SimpleNamespace(model=DiTConfig()) - self.device = torch.device("cpu") - self.diffusion_steps = config.diffusion_steps - self.model: DiTPolicy | None = None - self.embedder: CLIPTextEmbedder | None = None - self.task_vec = torch.empty(0) - self.norm_stats: dict[str, Any] = {} - self._task = "" - self._fast_graph: FastInferenceGraph | None = None - - def load(self, profile: PolicyIOProfile) -> PolicyBackendInfo: - _validate_profile(profile, self.config.model) - checkpoint = Path(self._rollout_config.artifact).expanduser().resolve() - if not checkpoint.is_file(): - raise FileNotFoundError(f"ABC checkpoint does not exist: {checkpoint}") - device_name = self._rollout_config.device or ( - "cuda" if torch.cuda.is_available() else "cpu" - ) - self.device = torch.device(device_name) - if self.device.type == "cuda" and not torch.cuda.is_available(): - raise RuntimeError(f"ABC requested device {device_name!r}, but CUDA is unavailable") - - self.model = DiTPolicy(self.config.model).to(self.device) - checkpoint_data = load_pretrained(self.model, checkpoint) - self.model.eval() - self.norm_stats = _resolve_norm_stats( - checkpoint_data, - self._rollout_config.norm_stats_path, - ) - _validate_norm_stats(self.norm_stats, self.config.model) - self.embedder = CLIPTextEmbedder(ClipConfig(), device=self.device) - self._set_task(self._rollout_config.task) - return PolicyBackendInfo( - name="abc", - chunk_length=self.config.model.chunk_length, - preferred_execution_steps=15, - ) - - def reset(self) -> None: - """ABC-DiT has no recurrent or queued model state.""" - - @torch.no_grad() - def predict( - self, - observations: Mapping[str, NDArray[Any]], - task: str, - ) -> NDArray[np.float32]: - model = self._require_model() - if task != self._task: - if self._fast_graph is not None: - raise ValueError("task text cannot change after ABC fast inference is captured") - self._set_task(task) - images: dict[str, NDArray[Any]] = { - camera: np.asarray(observations[camera]).transpose(2, 0, 1) - for camera in self.config.model.camera_keys - } - obs: dict[str, Any] = { - "state": np.asarray(observations["state"], dtype=np.float32), - "images": images, - } - if self._rollout_config.fast_inference and self.device.type == "cuda": - if self._fast_graph is None: - self._enable_fast_inference(obs) - assert self._fast_graph is not None - return self._fast_graph.infer(obs, noise=None) - - state = normalize(obs["state"], self.norm_stats["state"]) - batch = { - "state": torch.from_numpy(state[None]).float().to(self.device), - "actions": torch.zeros( - 1, - self.config.model.chunk_length, - self.config.model.action_dim, - device=self.device, - ), - "images": { - camera: resize_pad_normalize(obs["images"][camera]).unsqueeze(0).to(self.device) - for camera in self.config.model.camera_keys - }, - "task_vec_clip": self.task_vec, - } - actions = model.sample_actions(batch, num_steps=self.diffusion_steps) - result = actions[0].float().detach().cpu().numpy() - return np.asarray( - unnormalize(result, self.norm_stats["actions"]), - dtype=np.float32, - ) - - def _set_task(self, task: str) -> None: - if self.embedder is None: - raise RuntimeError("ABC text embedder is not loaded") - self.task_vec = self.embedder.encode([task]).to(self.device) - self._task = task - - def _enable_fast_inference(self, observation: dict[str, Any]) -> None: - model = self._require_model() - torch.backends.cuda.matmul.allow_tf32 = True - torch.backends.cudnn.allow_tf32 = True - model.to(torch.bfloat16) - model.img_backbone.set_bfloat16(True) - self.task_vec = self.task_vec.to(device=self.device, dtype=torch.bfloat16) - model.predict_velocity = torch.compile( # type: ignore[method-assign] - model.predict_velocity, - dynamic=False, - mode="max-autotune-no-cudagraphs", - ) - graph = FastInferenceGraph(self) - graph.capture(observation, warmup_noise=None, replay_warmups=24) - self._fast_graph = graph - - def _require_model(self) -> DiTPolicy: - if self.model is None: - raise RuntimeError("ABC backend is not loaded") - return self.model - - -AbcPolicyRuntime = declare_policy_runtime( - "AbcPolicyRuntime", - __name__, - DualOpenYamAbcPolicy, - AbcBackend, -) - - -def _validate_profile(profile: PolicyIOProfile, model: DiTConfig) -> None: - expected = {*model.camera_keys, "state"} - if set(profile.observations) != expected: - raise ValueError(f"released ABC checkpoint requires observation keys {sorted(expected)}") - for key in model.camera_keys: - if not isinstance(profile.observations[key], ImageSource): - raise TypeError(f"ABC observation {key!r} must be an image") - state = profile.observations["state"] - if not isinstance(state, JointPositionSource) or len(state.joints) != model.state_dim: - raise ValueError(f"ABC state must contain {model.state_dim} joints") - if profile.action.key != "actions": - raise ValueError("released ABC checkpoint requires the 'actions' output key") - if len(profile.action.demonstration.joints) != model.action_dim: - raise ValueError(f"ABC actions must contain {model.action_dim} joints") - - -def _resolve_norm_stats( - checkpoint: dict[str, Any], - override: str | None, -) -> dict[str, Any]: - if override is not None: - raw = json.loads(Path(override).expanduser().read_text()) - elif checkpoint.get("norm_stats") is not None: - raw = checkpoint["norm_stats"] - else: - raise ValueError("ABC checkpoint has no norm_stats; set norm_stats_path") - return cast("dict[str, Any]", parse_norm_stats(raw)) - - -def _validate_norm_stats(stats: dict[str, Any], model: DiTConfig) -> None: - for key, width in (("state", model.state_dim), ("actions", model.action_dim)): - for statistic in ("mean", "std"): - value = np.asarray(stats[key][statistic]) - if value.shape != (width,): - raise ValueError( - f"ABC {key} {statistic} shape {value.shape} does not match {(width,)}" - ) - if not np.all(np.isfinite(value)): - raise ValueError(f"ABC {key} {statistic} contains non-finite values") diff --git a/dimos/imitation/policy/abc/python/dimos_abc/runtime_tests.py b/dimos/imitation/policy/abc/python/dimos_abc/runtime_tests.py deleted file mode 100644 index 0d4ced30be..0000000000 --- a/dimos/imitation/policy/abc/python/dimos_abc/runtime_tests.py +++ /dev/null @@ -1,40 +0,0 @@ -# Copyright 2026 Dimensional Inc. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -from abc_minimal.config import DiTConfig -import numpy as np -import pytest - -from dimos.experimental.isolated_python.bootstrap import validate_runtime -from dimos.imitation.policy.abc.module import DualOpenYamAbcPolicy -from dimos.robot.manipulators.dual_openyam.learning import DUAL_OPENYAM_ABC_IO -from dimos_abc.runtime import AbcPolicyRuntime, _validate_norm_stats, _validate_profile - - -def test_generated_runtime_implements_the_host_contract() -> None: - validate_runtime(DualOpenYamAbcPolicy, AbcPolicyRuntime) - - -def test_released_abc_profile_matches_vendored_model_contract() -> None: - _validate_profile(DUAL_OPENYAM_ABC_IO, DiTConfig()) - - -def test_norm_stats_require_exact_released_dimensions() -> None: - stats = { - "state": {"mean": np.zeros(14), "std": np.ones(14)}, - "actions": {"mean": np.zeros(13), "std": np.ones(13)}, - } - - with pytest.raises(ValueError, match="actions mean shape"): - _validate_norm_stats(stats, DiTConfig()) diff --git a/dimos/imitation/policy/abc/python/pyproject.toml b/dimos/imitation/policy/abc/python/pyproject.toml deleted file mode 100644 index 62f4f57e07..0000000000 --- a/dimos/imitation/policy/abc/python/pyproject.toml +++ /dev/null @@ -1,56 +0,0 @@ -[build-system] -requires = ["setuptools>=70"] -build-backend = "setuptools.build_meta" - -[project] -name = "dimos-abc-runtime" -version = "0.1.0" -requires-python = ">=3.12,<3.13" -dependencies = [ - "ftfy>=6.3,<7", - "numpy>=2,<3", - "regex>=2025.7", - "torch==2.11.0+cu128", -] - -[[tool.uv.index]] -name = "pytorch-cu128" -url = "https://download.pytorch.org/whl/cu128" -explicit = true - -[tool.uv.sources] -torch = { index = "pytorch-cu128" } - -[dependency-groups] -tests = [ - "mypy==1.19.0", - "pytest==8.3.5", - "pytest-mock>=3.14", -] - -[tool.uv] -default-groups = [] - -[tool.setuptools.packages.find] -where = ["."] -include = ["abc_minimal*", "dimos_abc*"] - -[tool.pytest.ini_options] -testpaths = ["dimos_abc"] -python_files = ["*_tests.py"] - -[tool.mypy] -files = ["dimos_abc/"] -python_version = "3.12" -strict = true -explicit_package_bases = true -mypy_path = "../../../../../" -untyped_calls_exclude = ["abc_minimal"] - -[[tool.mypy.overrides]] -module = ["dimos", "dimos.*"] -follow_imports = "skip" - -[[tool.mypy.overrides]] -module = ["abc_minimal", "abc_minimal.*"] -follow_untyped_imports = true diff --git a/dimos/imitation/policy/abc/python/uv.lock b/dimos/imitation/policy/abc/python/uv.lock deleted file mode 100644 index 3ddbaf292e..0000000000 --- a/dimos/imitation/policy/abc/python/uv.lock +++ /dev/null @@ -1,569 +0,0 @@ -version = 1 -revision = 3 -requires-python = "==3.12.*" - -[[package]] -name = "colorama" -version = "0.4.6" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/d8/53/6f443c9a4a8358a93a6792e2acffb9d9d5cb0a5cfd8802644b7b1c9a02e4/colorama-0.4.6.tar.gz", hash = "sha256:08695f5cb7ed6e0531a20572697297273c47b8cae5a63ffc6d6ed5c201be6e44", size = 27697, upload-time = "2022-10-25T02:36:22.414Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/d1/d6/3965ed04c63042e047cb6a3e6ed1a63a35087b6a609aa3a15ed8ac56c221/colorama-0.4.6-py2.py3-none-any.whl", hash = "sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6", size = 25335, upload-time = "2022-10-25T02:36:20.889Z" }, -] - -[[package]] -name = "cuda-bindings" -version = "12.9.7" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "cuda-pathfinder" }, -] -wheels = [ - { url = "https://files.pythonhosted.org/packages/32/45/557d4ed1fa54f0c7db8aee083229f624990d69f7d00f55477eed5c7e169a/cuda_bindings-12.9.7-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:0666d3c082ef8f4b2d670950589373550e9f3bf564d635dd883f24a0b40402ff", size = 7071026, upload-time = "2026-05-27T18:44:13.356Z" }, - { url = "https://files.pythonhosted.org/packages/91/97/e3c6e58ece26a053419ba0a18444b5443cfc64451bbf37f84e8143b8bdca/cuda_bindings-12.9.7-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:4c7ef48c5e13ae90f3b2ecfb72f8e99ac43c8f4c43e67e1325b8aae331453687", size = 7611059, upload-time = "2026-05-27T18:44:15.252Z" }, -] - -[[package]] -name = "cuda-pathfinder" -version = "1.8.1" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/9d/e6/22df83f82f9bc26cb1c42265cf14d34d4908dba2a0f261bd7b28244acb00/cuda_pathfinder-1.8.1-py3-none-any.whl", hash = "sha256:ae0137ff9e56ea97499bcbf54f5f2778ec25f3266715ac86da192a795af982a8", size = 62552, upload-time = "2026-09-02T16:55:28.64Z" }, -] - -[[package]] -name = "cuda-toolkit" -version = "12.8.1" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/d4/c8/7dce3a0b15b42a3b58e7d96eb22a687d3bf2c44e01d149a6874629cd9938/cuda_toolkit-12.8.1-py2.py3-none-any.whl", hash = "sha256:adc7906af4ecbf9a352f9dca5734eceb21daec281ccfcf5675e1d2f724fc2cba", size = 2283, upload-time = "2025-08-13T02:03:07.842Z" }, -] - -[package.optional-dependencies] -cublas = [ - { name = "nvidia-cublas-cu12", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, -] -cudart = [ - { name = "nvidia-cuda-runtime-cu12", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, -] -cufft = [ - { name = "nvidia-cufft-cu12", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, -] -cufile = [ - { name = "nvidia-cufile-cu12", marker = "sys_platform == 'linux'" }, -] -cupti = [ - { name = "nvidia-cuda-cupti-cu12", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, -] -curand = [ - { name = "nvidia-curand-cu12", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, -] -cusolver = [ - { name = "nvidia-cusolver-cu12", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, -] -cusparse = [ - { name = "nvidia-cusparse-cu12", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, -] -nvjitlink = [ - { name = "nvidia-nvjitlink-cu12", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, -] -nvrtc = [ - { name = "nvidia-cuda-nvrtc-cu12", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, -] -nvtx = [ - { name = "nvidia-nvtx-cu12", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, -] - -[[package]] -name = "dimos-abc-runtime" -version = "0.1.0" -source = { editable = "." } -dependencies = [ - { name = "ftfy" }, - { name = "numpy" }, - { name = "regex" }, - { name = "torch" }, -] - -[package.dev-dependencies] -tests = [ - { name = "mypy" }, - { name = "pytest" }, - { name = "pytest-mock" }, -] - -[package.metadata] -requires-dist = [ - { name = "ftfy", specifier = ">=6.3,<7" }, - { name = "numpy", specifier = ">=2,<3" }, - { name = "regex", specifier = ">=2025.7" }, - { name = "torch", specifier = "==2.11.0+cu128", index = "https://download.pytorch.org/whl/cu128" }, -] - -[package.metadata.requires-dev] -tests = [ - { name = "mypy", specifier = "==1.19.0" }, - { name = "pytest", specifier = "==8.3.5" }, - { name = "pytest-mock", specifier = ">=3.14" }, -] - -[[package]] -name = "filelock" -version = "3.32.5" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/0a/a0/50c2c0ce5e74d7721bbb1b19a26ebd339aac5878553a6e35308c2f31f935/filelock-3.32.5.tar.gz", hash = "sha256:f6a6a28f743f9b95ce19db5abe0f376f75eb56517dff21e1a4751e2657d3e83d", size = 222838, upload-time = "2026-08-31T18:56:34.729Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/36/d2/b70a31e13d04456d28493f31d2aa087e99eeb2767ef0293b2625727ccb8c/filelock-3.32.5-py3-none-any.whl", hash = "sha256:142cd9fa77a872c5e78c62329a0d15278fadc686eb89e760017968961a4fd6b2", size = 100003, upload-time = "2026-08-31T18:56:33.078Z" }, -] - -[[package]] -name = "fsspec" -version = "2026.7.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/00/78/f34251dadb8f3921264a1d9b8946f5e542014ee2614b285261b4e40e6775/fsspec-2026.7.0.tar.gz", hash = "sha256:c803c40f4cf860b49dea58ee3e1c33cb9c790520e233537e1340049f89b82a88", size = 317040, upload-time = "2026-07-28T16:34:51.052Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/fd/3c/6a2bf344106328fd04963664a60b9bb6496fc25df8e962fcdc1367285fb9/fsspec-2026.7.0-py3-none-any.whl", hash = "sha256:b57ddbafedfaef7018c1ecab32aa200a9d7ca26b77965f64e48b70061249d279", size = 206583, upload-time = "2026-07-28T16:34:49.538Z" }, -] - -[[package]] -name = "ftfy" -version = "6.3.1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "wcwidth" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/a5/d3/8650919bc3c7c6e90ee3fa7fd618bf373cbbe55dff043bd67353dbb20cd8/ftfy-6.3.1.tar.gz", hash = "sha256:9b3c3d90f84fb267fe64d375a07b7f8912d817cf86009ae134aa03e1819506ec", size = 308927, upload-time = "2024-10-26T00:50:35.149Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/ab/6e/81d47999aebc1b155f81eca4477a616a70f238a2549848c38983f3c22a82/ftfy-6.3.1-py3-none-any.whl", hash = "sha256:7c70eb532015cd2f9adb53f101fb6c7945988d023a085d127d1573dc49dd0083", size = 44821, upload-time = "2024-10-26T00:50:33.425Z" }, -] - -[[package]] -name = "iniconfig" -version = "2.3.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/72/34/14ca021ce8e5dfedc35312d08ba8bf51fdd999c576889fc2c24cb97f4f10/iniconfig-2.3.0.tar.gz", hash = "sha256:c76315c77db068650d49c5b56314774a7804df16fee4402c1f19d6d15d8c4730", size = 20503, upload-time = "2025-10-18T21:55:43.219Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/cb/b1/3846dd7f199d53cb17f49cba7e651e9ce294d8497c8c150530ed11865bb8/iniconfig-2.3.0-py3-none-any.whl", hash = "sha256:f631c04d2c48c52b84d0d0549c99ff3859c98df65b3101406327ecc7d53fbf12", size = 7484, upload-time = "2025-10-18T21:55:41.639Z" }, -] - -[[package]] -name = "jinja2" -version = "3.1.6" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "markupsafe" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/df/bf/f7da0350254c0ed7c72f3e33cef02e048281fec7ecec5f032d4aac52226b/jinja2-3.1.6.tar.gz", hash = "sha256:0137fb05990d35f1275a587e9aee6d56da821fc83491a0fb838183be43f66d6d", size = 245115, upload-time = "2025-03-05T20:05:02.478Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/62/a1/3d680cbfd5f4b8f15abc1d571870c5fc3e594bb582bc3b64ea099db13e56/jinja2-3.1.6-py3-none-any.whl", hash = "sha256:85ece4451f492d0c13c5dd7c13a64681a86afae63a5f347908daf103ce6d2f67", size = 134899, upload-time = "2025-03-05T20:05:00.369Z" }, -] - -[[package]] -name = "librt" -version = "0.15.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/36/9b/356320fbae2ac8467e21c5e73e1389c80468e4998c62cc7d3536cc51b614/librt-0.15.0.tar.gz", hash = "sha256:4e66cbe84437497d951b799d3e1551291b6fb3d643820a7014b3655d57a59162", size = 214338, upload-time = "2026-08-07T10:49:42.663Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/ba/39/99c25030e782bdfb7a21be8c05254806a2e4bbb05c8d50c2a2130acbfa05/librt-0.15.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:e87bc679f86a99aa3b26e3c78eeb821a247c9a28eae48eaafcc32c3bf4c3bb9e", size = 151021, upload-time = "2026-08-07T10:47:00.057Z" }, - { url = "https://files.pythonhosted.org/packages/14/43/f4b1bd1b2888798a1409808889a25ea1ba49eaabce7d681ed27734c2df9d/librt-0.15.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:71599e011ac880e8e45d46047d714871894c7d4ab6f25626f8d4f89da21f368d", size = 155267, upload-time = "2026-08-07T10:47:01.311Z" }, - { url = "https://files.pythonhosted.org/packages/0c/db/3ad9c965c72f1e1d6beeec44ec10a54e17be8ae042fbb4baade16cbadced/librt-0.15.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c802434092b769b1d613ed2e13fac15fbfce1934a74bd10283b03c0fae231cd1", size = 503136, upload-time = "2026-08-07T10:47:02.45Z" }, - { url = "https://files.pythonhosted.org/packages/4b/07/5888a6d76acd62ebce66c61b74d94e9370b9c32929f111e487bb6546f8ed/librt-0.15.0-cp312-cp312-manylinux2014_i686.manylinux_2_17_i686.manylinux_2_28_i686.whl", hash = "sha256:5500eeae393a184d14e1f35645962c27129d20c81afa4069e6ef826ebc2b3aaa", size = 496670, upload-time = "2026-08-07T10:47:03.675Z" }, - { url = "https://files.pythonhosted.org/packages/29/39/ab57cc2f5b276156da02bb7f5a8921bada1cb1993ffec99acf811c602c23/librt-0.15.0-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:6ecfc32dfb46fb7b565bcd6abf9412acf978775a998273d22888a6d7953730dd", size = 513688, upload-time = "2026-08-07T10:47:04.981Z" }, - { url = "https://files.pythonhosted.org/packages/a7/b9/bdbb0b648b5c2befb031f4c6f3b1dd857415e8fb492a25a3c764a6681e6c/librt-0.15.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:89cc46cfd15022e35084355478c9ac809d90b1152222706ac9a7655ec21df6fa", size = 531904, upload-time = "2026-08-07T10:47:06.211Z" }, - { url = "https://files.pythonhosted.org/packages/93/26/473c2e4b6c104e9e58e27ce95fc8005c8bd4fc36cae4f254371125a92db8/librt-0.15.0-cp312-cp312-manylinux_2_34_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:d5f51401d102c885b9ca509e62c79b1dbff286e1b9b047fde6f763780789356d", size = 524427, upload-time = "2026-08-07T10:47:07.592Z" }, - { url = "https://files.pythonhosted.org/packages/26/60/03b3abb82b41714671b907bf6989b228e31e6a8af52dec82b5b0728dc250/librt-0.15.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:cc30523e3f1a23fb7511cc659834a0d01a1042bb9de359bc1c131cc4ec6c9656", size = 543155, upload-time = "2026-08-07T10:47:08.866Z" }, - { url = "https://files.pythonhosted.org/packages/f2/0e/9bb1f0a4affbd0a1888f4f79dc03ed2a299d9a2c26c59ab2a97dcbf11903/librt-0.15.0-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:59fe030d8ae4a57e3fb7756bf35a858de74e04066fc8555c53d0af979132af81", size = 546890, upload-time = "2026-08-07T10:47:10.327Z" }, - { url = "https://files.pythonhosted.org/packages/dc/84/6937a280d461f7de6e031ffb02edc2b7c3c90d49d630565ce8ff27cbc5f2/librt-0.15.0-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:5a6526a2a956bbb1e4ae3568c82e650fc99119c66bb011ea60715744955a2b4d", size = 555163, upload-time = "2026-08-07T10:47:11.798Z" }, - { url = "https://files.pythonhosted.org/packages/bc/95/2a2853c1ee014bf102116e7f897a04beeaeb2461b45b79af98bdfb95f1ef/librt-0.15.0-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:85ea21ec6730194d67156b0e0b5430ccb1d61f8b8b907e39b37f9812b74a13f0", size = 535812, upload-time = "2026-08-07T10:47:13.279Z" }, - { url = "https://files.pythonhosted.org/packages/c9/4c/cf9601c1b4c5f09280acd5d83abdb2e68527a2be8257136eb42304218622/librt-0.15.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:1e47b8ba865d7ede071a91a7163073bbaeb72541f1ef8a07d512c45c7b5007f2", size = 573688, upload-time = "2026-08-07T10:47:14.727Z" }, - { url = "https://files.pythonhosted.org/packages/47/6d/9ac7cbec46189a7625af4b5acbd25f10d827f4141b2002181848c8418923/librt-0.15.0-cp312-cp312-win32.whl", hash = "sha256:a5207ec414d1c4a2a7231b2086970dc036f94293cdf338190984958a013a42f1", size = 106138, upload-time = "2026-08-07T10:47:15.973Z" }, - { url = "https://files.pythonhosted.org/packages/38/d0/2ae99c83be86ce23f925ac1aeeedc777e97f427c4a8d190c70d0a16e9a87/librt-0.15.0-cp312-cp312-win_amd64.whl", hash = "sha256:73b30cfa976659b3917c8f6153bdb0591c6a9ec6583599fd24a689b690622022", size = 126974, upload-time = "2026-08-07T10:47:17.049Z" }, - { url = "https://files.pythonhosted.org/packages/5d/ef/dd24f9635c730b86b87587967dda7516b1845e8b17684603d31607fed598/librt-0.15.0-cp312-cp312-win_arm64.whl", hash = "sha256:a54cf9e0ef47b96af580849db5471142200568ce1e02cbf416addab551369570", size = 112292, upload-time = "2026-08-07T10:47:18.222Z" }, -] - -[[package]] -name = "markupsafe" -version = "3.0.3" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/7e/99/7690b6d4034fffd95959cbe0c02de8deb3098cc577c67bb6a24fe5d7caa7/markupsafe-3.0.3.tar.gz", hash = "sha256:722695808f4b6457b320fdc131280796bdceb04ab50fe1795cd540799ebe1698", size = 80313, upload-time = "2025-09-27T18:37:40.426Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/5a/72/147da192e38635ada20e0a2e1a51cf8823d2119ce8883f7053879c2199b5/markupsafe-3.0.3-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:d53197da72cc091b024dd97249dfc7794d6a56530370992a5e1a08983ad9230e", size = 11615, upload-time = "2025-09-27T18:36:30.854Z" }, - { url = "https://files.pythonhosted.org/packages/9a/81/7e4e08678a1f98521201c3079f77db69fb552acd56067661f8c2f534a718/markupsafe-3.0.3-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:1872df69a4de6aead3491198eaf13810b565bdbeec3ae2dc8780f14458ec73ce", size = 12020, upload-time = "2025-09-27T18:36:31.971Z" }, - { url = "https://files.pythonhosted.org/packages/1e/2c/799f4742efc39633a1b54a92eec4082e4f815314869865d876824c257c1e/markupsafe-3.0.3-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:3a7e8ae81ae39e62a41ec302f972ba6ae23a5c5396c8e60113e9066ef893da0d", size = 24332, upload-time = "2025-09-27T18:36:32.813Z" }, - { url = "https://files.pythonhosted.org/packages/3c/2e/8d0c2ab90a8c1d9a24f0399058ab8519a3279d1bd4289511d74e909f060e/markupsafe-3.0.3-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d6dd0be5b5b189d31db7cda48b91d7e0a9795f31430b7f271219ab30f1d3ac9d", size = 22947, upload-time = "2025-09-27T18:36:33.86Z" }, - { url = "https://files.pythonhosted.org/packages/2c/54/887f3092a85238093a0b2154bd629c89444f395618842e8b0c41783898ea/markupsafe-3.0.3-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:94c6f0bb423f739146aec64595853541634bde58b2135f27f61c1ffd1cd4d16a", size = 21962, upload-time = "2025-09-27T18:36:35.099Z" }, - { url = "https://files.pythonhosted.org/packages/c9/2f/336b8c7b6f4a4d95e91119dc8521402461b74a485558d8f238a68312f11c/markupsafe-3.0.3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:be8813b57049a7dc738189df53d69395eba14fb99345e0a5994914a3864c8a4b", size = 23760, upload-time = "2025-09-27T18:36:36.001Z" }, - { url = "https://files.pythonhosted.org/packages/32/43/67935f2b7e4982ffb50a4d169b724d74b62a3964bc1a9a527f5ac4f1ee2b/markupsafe-3.0.3-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:83891d0e9fb81a825d9a6d61e3f07550ca70a076484292a70fde82c4b807286f", size = 21529, upload-time = "2025-09-27T18:36:36.906Z" }, - { url = "https://files.pythonhosted.org/packages/89/e0/4486f11e51bbba8b0c041098859e869e304d1c261e59244baa3d295d47b7/markupsafe-3.0.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:77f0643abe7495da77fb436f50f8dab76dbc6e5fd25d39589a0f1fe6548bfa2b", size = 23015, upload-time = "2025-09-27T18:36:37.868Z" }, - { url = "https://files.pythonhosted.org/packages/2f/e1/78ee7a023dac597a5825441ebd17170785a9dab23de95d2c7508ade94e0e/markupsafe-3.0.3-cp312-cp312-win32.whl", hash = "sha256:d88b440e37a16e651bda4c7c2b930eb586fd15ca7406cb39e211fcff3bf3017d", size = 14540, upload-time = "2025-09-27T18:36:38.761Z" }, - { url = "https://files.pythonhosted.org/packages/aa/5b/bec5aa9bbbb2c946ca2733ef9c4ca91c91b6a24580193e891b5f7dbe8e1e/markupsafe-3.0.3-cp312-cp312-win_amd64.whl", hash = "sha256:26a5784ded40c9e318cfc2bdb30fe164bdb8665ded9cd64d500a34fb42067b1c", size = 15105, upload-time = "2025-09-27T18:36:39.701Z" }, - { url = "https://files.pythonhosted.org/packages/e5/f1/216fc1bbfd74011693a4fd837e7026152e89c4bcf3e77b6692fba9923123/markupsafe-3.0.3-cp312-cp312-win_arm64.whl", hash = "sha256:35add3b638a5d900e807944a078b51922212fb3dedb01633a8defc4b01a3c85f", size = 13906, upload-time = "2025-09-27T18:36:40.689Z" }, -] - -[[package]] -name = "mpmath" -version = "1.3.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/e0/47/dd32fa426cc72114383ac549964eecb20ecfd886d1e5ccf5340b55b02f57/mpmath-1.3.0.tar.gz", hash = "sha256:7a28eb2a9774d00c7bc92411c19a89209d5da7c4c9a9e227be8330a23a25b91f", size = 508106, upload-time = "2023-03-07T16:47:11.061Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/43/e3/7d92a15f894aa0c9c4b49b8ee9ac9850d6e63b03c9c32c0367a13ae62209/mpmath-1.3.0-py3-none-any.whl", hash = "sha256:a0b2b9fe80bbcd81a6647ff13108738cfb482d481d826cc0e02f5b35e5c88d2c", size = 536198, upload-time = "2023-03-07T16:47:09.197Z" }, -] - -[[package]] -name = "mypy" -version = "1.19.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "librt" }, - { name = "mypy-extensions" }, - { name = "pathspec" }, - { name = "typing-extensions" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/f9/b5/b58cdc25fadd424552804bf410855d52324183112aa004f0732c5f6324cf/mypy-1.19.0.tar.gz", hash = "sha256:f6b874ca77f733222641e5c46e4711648c4037ea13646fd0cdc814c2eaec2528", size = 3579025, upload-time = "2025-11-28T15:49:01.26Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/11/7e/1afa8fb188b876abeaa14460dc4983f909aaacaa4bf5718c00b2c7e0b3d5/mypy-1.19.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:0fb3115cb8fa7c5f887c8a8d81ccdcb94cff334684980d847e5a62e926910e1d", size = 13207728, upload-time = "2025-11-28T15:46:26.463Z" }, - { url = "https://files.pythonhosted.org/packages/b2/13/f103d04962bcbefb1644f5ccb235998b32c337d6c13145ea390b9da47f3e/mypy-1.19.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:f3e19e3b897562276bb331074d64c076dbdd3e79213f36eed4e592272dabd760", size = 12202945, upload-time = "2025-11-28T15:48:49.143Z" }, - { url = "https://files.pythonhosted.org/packages/e4/93/a86a5608f74a22284a8ccea8592f6e270b61f95b8588951110ad797c2ddd/mypy-1.19.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:b9d491295825182fba01b6ffe2c6fe4e5a49dbf4e2bb4d1217b6ced3b4797bc6", size = 12718673, upload-time = "2025-11-28T15:47:37.193Z" }, - { url = "https://files.pythonhosted.org/packages/3d/58/cf08fff9ced0423b858f2a7495001fda28dc058136818ee9dffc31534ea9/mypy-1.19.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:6016c52ab209919b46169651b362068f632efcd5eb8ef9d1735f6f86da7853b2", size = 13608336, upload-time = "2025-11-28T15:48:32.625Z" }, - { url = "https://files.pythonhosted.org/packages/64/ed/9c509105c5a6d4b73bb08733102a3ea62c25bc02c51bca85e3134bf912d3/mypy-1.19.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:f188dcf16483b3e59f9278c4ed939ec0254aa8a60e8fc100648d9ab5ee95a431", size = 13833174, upload-time = "2025-11-28T15:45:48.091Z" }, - { url = "https://files.pythonhosted.org/packages/cd/71/01939b66e35c6f8cb3e6fdf0b657f0fd24de2f8ba5e523625c8e72328208/mypy-1.19.0-cp312-cp312-win_amd64.whl", hash = "sha256:0e3c3d1e1d62e678c339e7ade72746a9e0325de42cd2cccc51616c7b2ed1a018", size = 10112208, upload-time = "2025-11-28T15:46:41.702Z" }, - { url = "https://files.pythonhosted.org/packages/09/0e/fe228ed5aeab470c6f4eb82481837fadb642a5aa95cc8215fd2214822c10/mypy-1.19.0-py3-none-any.whl", hash = "sha256:0c01c99d626380752e527d5ce8e69ffbba2046eb8a060db0329690849cf9b6f9", size = 2469714, upload-time = "2025-11-28T15:45:33.22Z" }, -] - -[[package]] -name = "mypy-extensions" -version = "1.1.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/a2/6e/371856a3fb9d31ca8dac321cda606860fa4548858c0cc45d9d1d4ca2628b/mypy_extensions-1.1.0.tar.gz", hash = "sha256:52e68efc3284861e772bbcd66823fde5ae21fd2fdb51c62a211403730b916558", size = 6343, upload-time = "2025-04-22T14:54:24.164Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/79/7b/2c79738432f5c924bef5071f933bcc9efd0473bac3b4aa584a6f7c1c8df8/mypy_extensions-1.1.0-py3-none-any.whl", hash = "sha256:1be4cccdb0f2482337c4743e60421de3a356cd97508abadd57d47403e94f5505", size = 4963, upload-time = "2025-04-22T14:54:22.983Z" }, -] - -[[package]] -name = "networkx" -version = "3.6.1" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/6a/51/63fe664f3908c97be9d2e4f1158eb633317598cfa6e1fc14af5383f17512/networkx-3.6.1.tar.gz", hash = "sha256:26b7c357accc0c8cde558ad486283728b65b6a95d85ee1cd66bafab4c8168509", size = 2517025, upload-time = "2025-12-08T17:02:39.908Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/9e/c9/b2622292ea83fbb4ec318f5b9ab867d0a28ab43c5717bb85b0a5f6b3b0a4/networkx-3.6.1-py3-none-any.whl", hash = "sha256:d47fbf302e7d9cbbb9e2555a0d267983d2aa476bac30e90dfbe5669bd57f3762", size = 2068504, upload-time = "2025-12-08T17:02:38.159Z" }, -] - -[[package]] -name = "numpy" -version = "2.5.2" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/9a/80/db0b4559e57ec36362bedbb05530a87fafbcb6067708c946967a41d449e7/numpy-2.5.2.tar.gz", hash = "sha256:d482d171c406ae88c5b19cad3b6a1c4c5209f886ab74bc44c2c865c23f52d860", size = 20773161, upload-time = "2026-08-09T13:48:27.962Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/69/72/dccb0aaf40972777283303919f613964227266d0c13adebb79ac124f1c3e/numpy-2.5.2-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:14e373cfc6387177e8409dac3c7159be8eb05cd77096cd7c950268b86f62831c", size = 16891693, upload-time = "2026-08-09T13:44:51.702Z" }, - { url = "https://files.pythonhosted.org/packages/60/2e/b5aee50a1f74ac815cf8331812cb8251e29024025de462e0c047641c614c/numpy-2.5.2-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:4bbd96c833ecc8cc069ce518078fc8c60cb9cbfb0fea5b7a803ad65035596d03", size = 11903109, upload-time = "2026-08-09T13:44:55.501Z" }, - { url = "https://files.pythonhosted.org/packages/f3/f4/29e78102a80601cf034d4e9767022cffeca2c3b4c926e1754572ca95593d/numpy-2.5.2-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:6e8172ddfcf5cf74b811d372b570b83c60bd2de87a6fbfbebdadb4a9bd9c6cbb", size = 5350202, upload-time = "2026-08-09T13:44:58.401Z" }, - { url = "https://files.pythonhosted.org/packages/11/4b/dcd3b7eadaf4035d2c7a4289d232523a6964f602598ef7674e4bd7291f93/numpy-2.5.2-cp312-cp312-macosx_14_0_x86_64.whl", hash = "sha256:65f188481f1669e26f62b701e8205d19e460fa4a9b52a1414ba382330e4a3414", size = 6687736, upload-time = "2026-08-09T13:45:00.813Z" }, - { url = "https://files.pythonhosted.org/packages/e5/21/4947e0e9d6c9fc2e2ff15b8949049ee44f63adb9cacc729ab8793f97e712/numpy-2.5.2-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:8ee9c4eeb8454b3660a8b53493563c3e121c2fc94fbd72b848ef814ed7b676a9", size = 15612696, upload-time = "2026-08-09T13:45:04.151Z" }, - { url = "https://files.pythonhosted.org/packages/3a/5f/62d28cf019460c7f1394105b4d49d9911a9c444cb77ab0bd95a204c5a6de/numpy-2.5.2-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:3cdec01fa790a186d430433fdd4d4ffb70eed6f0eeb4bf05c8dbe2dce0a9bcb8", size = 16722264, upload-time = "2026-08-09T13:45:07.714Z" }, - { url = "https://files.pythonhosted.org/packages/14/25/3f0be4c1b9fdf5dd5e708a6806978564d7c46a055c000496309ff2a2f8af/numpy-2.5.2-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:7999d4ddb0c4025018373fd787510d46e04c769467af22869707b3c1cfd459ab", size = 16974396, upload-time = "2026-08-09T13:45:11.316Z" }, - { url = "https://files.pythonhosted.org/packages/22/72/6262cbdeeb45da9d971e40715f579d791603ba8ec0b5e2db1ac55454421d/numpy-2.5.2-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:c1f017dc0875c9209d219f97feceb7d54c2661bb243deb4114478e1295808af7", size = 18476044, upload-time = "2026-08-09T13:45:14.869Z" }, - { url = "https://files.pythonhosted.org/packages/36/33/29208b8b075bde62d26a81d14b358c42b0f69b6cabd98d4ff97f37f22b05/numpy-2.5.2-cp312-cp312-win32.whl", hash = "sha256:d6a48072864e3324e194a8fbb3c657bcc5b5c869dbc64c9537b1d5c862572c0a", size = 6072817, upload-time = "2026-08-09T13:45:17.867Z" }, - { url = "https://files.pythonhosted.org/packages/7f/b9/87fea2769fe1c47c1b5b01d8310772c9d1a85d485de7cf386ef7a3332b02/numpy-2.5.2-cp312-cp312-win_amd64.whl", hash = "sha256:28ac63476ec7651484215ee7fa15a1f78b57c14621f01e392afe17b9a1390ce4", size = 12464674, upload-time = "2026-08-09T13:45:20.734Z" }, - { url = "https://files.pythonhosted.org/packages/14/52/032b97e00461ab0809bbe4c588b035620e5a14b8cdee47ecddefc7b17d33/numpy-2.5.2-cp312-cp312-win_arm64.whl", hash = "sha256:27650bb0e7140fa3d37b9923b4803645e0b125d190f326eecfd3f4dad8e8ade1", size = 10397131, upload-time = "2026-08-09T13:45:23.73Z" }, -] - -[[package]] -name = "nvidia-cublas-cu12" -version = "12.8.4.1" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/29/99/db44d685f0e257ff0e213ade1964fc459b4a690a73293220e98feb3307cf/nvidia_cublas_cu12-12.8.4.1-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:b86f6dd8935884615a0683b663891d43781b819ac4f2ba2b0c9604676af346d0", size = 590537124, upload-time = "2025-03-07T01:43:53.556Z" }, - { url = "https://files.pythonhosted.org/packages/dc/61/e24b560ab2e2eaeb3c839129175fb330dfcfc29e5203196e5541a4c44682/nvidia_cublas_cu12-12.8.4.1-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:8ac4e771d5a348c551b2a426eda6193c19aa630236b418086020df5ba9667142", size = 594346921, upload-time = "2025-03-07T01:44:31.254Z" }, -] - -[[package]] -name = "nvidia-cuda-cupti-cu12" -version = "12.8.90" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/d5/1f/b3bd73445e5cb342727fd24fe1f7b748f690b460acadc27ea22f904502c8/nvidia_cuda_cupti_cu12-12.8.90-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:4412396548808ddfed3f17a467b104ba7751e6b58678a4b840675c56d21cf7ed", size = 9533318, upload-time = "2025-03-07T01:40:10.421Z" }, - { url = "https://files.pythonhosted.org/packages/f8/02/2adcaa145158bf1a8295d83591d22e4103dbfd821bcaf6f3f53151ca4ffa/nvidia_cuda_cupti_cu12-12.8.90-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:ea0cb07ebda26bb9b29ba82cda34849e73c166c18162d3913575b0c9db9a6182", size = 10248621, upload-time = "2025-03-07T01:40:21.213Z" }, -] - -[[package]] -name = "nvidia-cuda-nvrtc-cu12" -version = "12.8.93" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/05/6b/32f747947df2da6994e999492ab306a903659555dddc0fbdeb9d71f75e52/nvidia_cuda_nvrtc_cu12-12.8.93-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:a7756528852ef889772a84c6cd89d41dfa74667e24cca16bb31f8f061e3e9994", size = 88040029, upload-time = "2025-03-07T01:42:13.562Z" }, - { url = "https://files.pythonhosted.org/packages/eb/d1/e50d0acaab360482034b84b6e27ee83c6738f7d32182b987f9c7a4e32962/nvidia_cuda_nvrtc_cu12-12.8.93-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:fc1fec1e1637854b4c0a65fb9a8346b51dd9ee69e61ebaccc82058441f15bce8", size = 43106076, upload-time = "2025-03-07T01:41:59.817Z" }, -] - -[[package]] -name = "nvidia-cuda-runtime-cu12" -version = "12.8.90" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/7c/75/f865a3b236e4647605ea34cc450900854ba123834a5f1598e160b9530c3a/nvidia_cuda_runtime_cu12-12.8.90-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:52bf7bbee900262ffefe5e9d5a2a69a30d97e2bc5bb6cc866688caa976966e3d", size = 965265, upload-time = "2025-03-07T01:39:43.533Z" }, - { url = "https://files.pythonhosted.org/packages/0d/9b/a997b638fcd068ad6e4d53b8551a7d30fe8b404d6f1804abf1df69838932/nvidia_cuda_runtime_cu12-12.8.90-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:adade8dcbd0edf427b7204d480d6066d33902cab2a4707dcfc48a2d0fd44ab90", size = 954765, upload-time = "2025-03-07T01:40:01.615Z" }, -] - -[[package]] -name = "nvidia-cudnn-cu12" -version = "9.19.0.56" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "nvidia-cublas-cu12" }, -] -wheels = [ - { url = "https://files.pythonhosted.org/packages/09/b8/277c51962ee46fa3e5b203ac5f76107c650f781d6891e681e28e6f3e9fe6/nvidia_cudnn_cu12-9.19.0.56-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:08caaf27fe556aca82a3ee3b5aa49a77e7de0cfcb7ff4e5c29da426387a8267e", size = 656910700, upload-time = "2026-02-03T20:40:25.508Z" }, - { url = "https://files.pythonhosted.org/packages/c5/41/65225d42fba06fb3dd3972485ea258e7dd07a40d6e01c95da6766ad87354/nvidia_cudnn_cu12-9.19.0.56-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:ac6ad90a075bb33a94f2b4cf4622eac13dd4dc65cf6dd9c7572a318516a36625", size = 657906812, upload-time = "2026-02-03T20:44:12.638Z" }, -] - -[[package]] -name = "nvidia-cufft-cu12" -version = "11.3.3.83" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "nvidia-nvjitlink-cu12" }, -] -wheels = [ - { url = "https://files.pythonhosted.org/packages/60/bc/7771846d3a0272026c416fbb7e5f4c1f146d6d80704534d0b187dd6f4800/nvidia_cufft_cu12-11.3.3.83-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:848ef7224d6305cdb2a4df928759dca7b1201874787083b6e7550dd6765ce69a", size = 193109211, upload-time = "2025-03-07T01:44:56.873Z" }, - { url = "https://files.pythonhosted.org/packages/1f/13/ee4e00f30e676b66ae65b4f08cb5bcbb8392c03f54f2d5413ea99a5d1c80/nvidia_cufft_cu12-11.3.3.83-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:4d2dd21ec0b88cf61b62e6b43564355e5222e4a3fb394cac0db101f2dd0d4f74", size = 193118695, upload-time = "2025-03-07T01:45:27.821Z" }, -] - -[[package]] -name = "nvidia-cufile-cu12" -version = "1.13.1.3" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/bb/fe/1bcba1dfbfb8d01be8d93f07bfc502c93fa23afa6fd5ab3fc7c1df71038a/nvidia_cufile_cu12-1.13.1.3-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:1d069003be650e131b21c932ec3d8969c1715379251f8d23a1860554b1cb24fc", size = 1197834, upload-time = "2025-03-07T01:45:50.723Z" }, - { url = "https://files.pythonhosted.org/packages/1e/f5/5607710447a6fe9fd9b3283956fceeee8a06cda1d2f56ce31371f595db2a/nvidia_cufile_cu12-1.13.1.3-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:4beb6d4cce47c1a0f1013d72e02b0994730359e17801d395bdcbf20cfb3bb00a", size = 1120705, upload-time = "2025-03-07T01:45:41.434Z" }, -] - -[[package]] -name = "nvidia-curand-cu12" -version = "10.3.9.90" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/45/5e/92aa15eca622a388b80fbf8375d4760738df6285b1e92c43d37390a33a9a/nvidia_curand_cu12-10.3.9.90-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:dfab99248034673b779bc6decafdc3404a8a6f502462201f2f31f11354204acd", size = 63625754, upload-time = "2025-03-07T01:46:10.735Z" }, - { url = "https://files.pythonhosted.org/packages/fb/aa/6584b56dc84ebe9cf93226a5cde4d99080c8e90ab40f0c27bda7a0f29aa1/nvidia_curand_cu12-10.3.9.90-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:b32331d4f4df5d6eefa0554c565b626c7216f87a06a4f56fab27c3b68a830ec9", size = 63619976, upload-time = "2025-03-07T01:46:23.323Z" }, -] - -[[package]] -name = "nvidia-cusolver-cu12" -version = "11.7.3.90" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "nvidia-cublas-cu12" }, - { name = "nvidia-cusparse-cu12" }, - { name = "nvidia-nvjitlink-cu12" }, -] -wheels = [ - { url = "https://files.pythonhosted.org/packages/c8/32/f7cd6ce8a7690544d084ea21c26e910a97e077c9b7f07bf5de623ee19981/nvidia_cusolver_cu12-11.7.3.90-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:db9ed69dbef9715071232caa9b69c52ac7de3a95773c2db65bdba85916e4e5c0", size = 267229841, upload-time = "2025-03-07T01:46:54.356Z" }, - { url = "https://files.pythonhosted.org/packages/85/48/9a13d2975803e8cf2777d5ed57b87a0b6ca2cc795f9a4f59796a910bfb80/nvidia_cusolver_cu12-11.7.3.90-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:4376c11ad263152bd50ea295c05370360776f8c3427b30991df774f9fb26c450", size = 267506905, upload-time = "2025-03-07T01:47:16.273Z" }, -] - -[[package]] -name = "nvidia-cusparse-cu12" -version = "12.5.8.93" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "nvidia-nvjitlink-cu12" }, -] -wheels = [ - { url = "https://files.pythonhosted.org/packages/bc/f7/cd777c4109681367721b00a106f491e0d0d15cfa1fd59672ce580ce42a97/nvidia_cusparse_cu12-12.5.8.93-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:9b6c161cb130be1a07a27ea6923df8141f3c295852f4b260c65f18f3e0a091dc", size = 288117129, upload-time = "2025-03-07T01:47:40.407Z" }, - { url = "https://files.pythonhosted.org/packages/c2/f5/e1854cb2f2bcd4280c44736c93550cc300ff4b8c95ebe370d0aa7d2b473d/nvidia_cusparse_cu12-12.5.8.93-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:1ec05d76bbbd8b61b06a80e1eaf8cf4959c3d4ce8e711b65ebd0443bb0ebb13b", size = 288216466, upload-time = "2025-03-07T01:48:13.779Z" }, -] - -[[package]] -name = "nvidia-cusparselt-cu12" -version = "0.7.1" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/73/b9/598f6ff36faaece4b3c50d26f50e38661499ff34346f00e057760b35cc9d/nvidia_cusparselt_cu12-0.7.1-py3-none-manylinux2014_aarch64.whl", hash = "sha256:8878dce784d0fac90131b6817b607e803c36e629ba34dc5b433471382196b6a5", size = 283835557, upload-time = "2025-02-26T00:16:54.265Z" }, - { url = "https://files.pythonhosted.org/packages/56/79/12978b96bd44274fe38b5dde5cfb660b1d114f70a65ef962bcbbed99b549/nvidia_cusparselt_cu12-0.7.1-py3-none-manylinux2014_x86_64.whl", hash = "sha256:f1bb701d6b930d5a7cea44c19ceb973311500847f81b634d802b7b539dc55623", size = 287193691, upload-time = "2025-02-26T00:15:44.104Z" }, -] - -[[package]] -name = "nvidia-nccl-cu12" -version = "2.28.9" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/08/c4/120d2dfd92dff2c776d68f361ff8705fdea2ca64e20b612fab0fd3f581ac/nvidia_nccl_cu12-2.28.9-py3-none-manylinux_2_18_aarch64.whl", hash = "sha256:50a36e01c4a090b9f9c47d92cec54964de6b9fcb3362d0e19b8ffc6323c21b60", size = 296766525, upload-time = "2025-11-18T05:49:16.094Z" }, - { url = "https://files.pythonhosted.org/packages/4a/4e/44dbb46b3d1b0ec61afda8e84837870f2f9ace33c564317d59b70bc19d3e/nvidia_nccl_cu12-2.28.9-py3-none-manylinux_2_18_x86_64.whl", hash = "sha256:485776daa8447da5da39681af455aa3b2c2586ddcf4af8772495e7c532c7e5ab", size = 296782137, upload-time = "2025-11-18T05:49:34.248Z" }, -] - -[[package]] -name = "nvidia-nvjitlink-cu12" -version = "12.8.93" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/f6/74/86a07f1d0f42998ca31312f998bd3b9a7eff7f52378f4f270c8679c77fb9/nvidia_nvjitlink_cu12-12.8.93-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:81ff63371a7ebd6e6451970684f916be2eab07321b73c9d244dc2b4da7f73b88", size = 39254836, upload-time = "2025-03-07T01:49:55.661Z" }, - { url = "https://files.pythonhosted.org/packages/2a/a2/8cee5da30d13430e87bf99bb33455d2724d0a4a9cb5d7926d80ccb96d008/nvidia_nvjitlink_cu12-12.8.93-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:adccd7161ace7261e01bb91e44e88da350895c270d23f744f0820c818b7229e7", size = 38386204, upload-time = "2025-03-07T01:49:43.612Z" }, -] - -[[package]] -name = "nvidia-nvshmem-cu12" -version = "3.4.5" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/1d/6a/03aa43cc9bd3ad91553a88b5f6fb25ed6a3752ae86ce2180221962bc2aa5/nvidia_nvshmem_cu12-3.4.5-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:0b48363fc6964dede448029434c6abed6c5e37f823cb43c3bcde7ecfc0457e15", size = 138936938, upload-time = "2025-09-06T00:32:05.589Z" }, - { url = "https://files.pythonhosted.org/packages/b5/09/6ea3ea725f82e1e76684f0708bbedd871fc96da89945adeba65c3835a64c/nvidia_nvshmem_cu12-3.4.5-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:042f2500f24c021db8a06c5eec2539027d57460e1c1a762055a6554f72c369bd", size = 139103095, upload-time = "2025-09-06T00:32:31.266Z" }, -] - -[[package]] -name = "nvidia-nvtx-cu12" -version = "12.8.90" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/10/c0/1b303feea90d296f6176f32a2a70b5ef230f9bdeb3a72bddb0dc922dc137/nvidia_nvtx_cu12-12.8.90-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:d7ad891da111ebafbf7e015d34879f7112832fc239ff0d7d776b6cb685274615", size = 91161, upload-time = "2025-03-07T01:42:23.922Z" }, - { url = "https://files.pythonhosted.org/packages/a2/eb/86626c1bbc2edb86323022371c39aa48df6fd8b0a1647bc274577f72e90b/nvidia_nvtx_cu12-12.8.90-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:5b17e2001cc0d751a5bc2c6ec6d26ad95913324a4adb86788c944f8ce9ba441f", size = 89954, upload-time = "2025-03-07T01:42:44.131Z" }, -] - -[[package]] -name = "packaging" -version = "26.3" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/7d/fa/3944b40b07da9ce895c0e6303a5ab7d53da063554f534556b134a54d6093/packaging-26.3.tar.gz", hash = "sha256:94edc256424af38762eb31306eed28beb9f0efc50a8837492c9d6fd6004aed79", size = 313412, upload-time = "2026-08-04T18:15:28.737Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/63/34/ba1c580383c9eada3711951fef0795c80b829a078d72188184bcab9dd527/packaging-26.3-py3-none-any.whl", hash = "sha256:d7193f7c8e4e93f444fde0262bf90af30e16fa0ad0ad44cb553c87339b23cd1c", size = 129956, upload-time = "2026-08-04T18:15:27.159Z" }, -] - -[[package]] -name = "pathspec" -version = "1.1.1" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/5a/82/42f767fc1c1143d6fd36efb827202a2d997a375e160a71eb2888a925aac1/pathspec-1.1.1.tar.gz", hash = "sha256:17db5ecd524104a120e173814c90367a96a98d07c45b2e10c2f3919fff91bf5a", size = 135180, upload-time = "2026-04-27T01:46:08.907Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/f1/d9/7fb5aa316bc299258e68c73ba3bddbc499654a07f151cba08f6153988714/pathspec-1.1.1-py3-none-any.whl", hash = "sha256:a00ce642f577bf7f473932318056212bc4f8bfdf53128c78bbd5af0b9b20b189", size = 57328, upload-time = "2026-04-27T01:46:07.06Z" }, -] - -[[package]] -name = "pluggy" -version = "1.6.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/f9/e2/3e91f31a7d2b083fe6ef3fa267035b518369d9511ffab804f839851d2779/pluggy-1.6.0.tar.gz", hash = "sha256:7dcc130b76258d33b90f61b658791dede3486c3e6bfb003ee5c9bfb396dd22f3", size = 69412, upload-time = "2025-05-15T12:30:07.975Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl", hash = "sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746", size = 20538, upload-time = "2025-05-15T12:30:06.134Z" }, -] - -[[package]] -name = "pytest" -version = "8.3.5" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "colorama", marker = "sys_platform == 'win32'" }, - { name = "iniconfig" }, - { name = "packaging" }, - { name = "pluggy" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/ae/3c/c9d525a414d506893f0cd8a8d0de7706446213181570cdbd766691164e40/pytest-8.3.5.tar.gz", hash = "sha256:f4efe70cc14e511565ac476b57c279e12a855b11f48f212af1080ef2263d3845", size = 1450891, upload-time = "2025-03-02T12:54:54.503Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/30/3d/64ad57c803f1fa1e963a7946b6e0fea4a70df53c1a7fed304586539c2bac/pytest-8.3.5-py3-none-any.whl", hash = "sha256:c69214aa47deac29fad6c2a4f590b9c4a9fdb16a403176fe154b79c0b4d4d820", size = 343634, upload-time = "2025-03-02T12:54:52.069Z" }, -] - -[[package]] -name = "pytest-mock" -version = "3.15.1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "pytest" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/68/14/eb014d26be205d38ad5ad20d9a80f7d201472e08167f0bb4361e251084a9/pytest_mock-3.15.1.tar.gz", hash = "sha256:1849a238f6f396da19762269de72cb1814ab44416fa73a8686deac10b0d87a0f", size = 34036, upload-time = "2025-09-16T16:37:27.081Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/5a/cc/06253936f4a7fa2e0f48dfe6d851d9c56df896a9ab09ac019d70b760619c/pytest_mock-3.15.1-py3-none-any.whl", hash = "sha256:0a25e2eb88fe5168d535041d09a4529a188176ae608a6d249ee65abc0949630d", size = 10095, upload-time = "2025-09-16T16:37:25.734Z" }, -] - -[[package]] -name = "regex" -version = "2026.9.3" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/19/c1/6b30b775c7bcc6cf6506a4d4741c2123e8d99cd50f3fe8cbd731f5fef526/regex-2026.9.3.tar.gz", hash = "sha256:aabd43208e335f4c3f0b56de3464b066dd425983a58f6eeb5738bcd7465403db", size = 416720, upload-time = "2026-09-01T00:53:43.821Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/da/cb/cba530bc3b068fc337f8f455c63ef5ee91a4eb4c76ecf5998e5cef5aaa6b/regex-2026.9.3-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:5db80d0b1c8238940b5957dd66b5c818ea40a221f6652fb717c027a562d09c77", size = 496699, upload-time = "2026-09-01T00:50:27.98Z" }, - { url = "https://files.pythonhosted.org/packages/81/39/f2e9fb6bbbc80f8bf67ad79d7e2e8866f7837d7c24c692f7faf8f1272e7e/regex-2026.9.3-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:35d48ce3dee087b63b15cd0a7a3110d0a76c29edbe1f2ad0520b8c4adb7cb596", size = 297018, upload-time = "2026-09-01T00:50:29.487Z" }, - { url = "https://files.pythonhosted.org/packages/a1/b6/c16ee58840baf7659def27ef6f62f3d9a9909670d3c1b4b98bb8b8ee47e2/regex-2026.9.3-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:1f22e0d21ae7016c77175c139a7fca465b988efc1280df4816c79752068d9e2e", size = 292008, upload-time = "2026-09-01T00:50:30.929Z" }, - { url = "https://files.pythonhosted.org/packages/cb/b4/4987bf0f17604669b4ea5aef219886d0a73188c4716ff3a7d275d4d15c15/regex-2026.9.3-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:233662cf8cfdfe3c0e58aa8f7bbefc579b5be0ac34546f123c159804179e8687", size = 796101, upload-time = "2026-09-01T00:50:32.486Z" }, - { url = "https://files.pythonhosted.org/packages/0b/95/2a9ab02a68c8a61dc0b4882ed643b1a95740d9dc291dc26c77d19af79691/regex-2026.9.3-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:7d2eed2e4d231278a2ccab3f4bfa2c1e39855f336475f7756a281d767d2b1753", size = 865435, upload-time = "2026-09-01T00:50:34.171Z" }, - { url = "https://files.pythonhosted.org/packages/04/92/0570d41559b446c97c1148cb9ebc1df09f2949b03c7c9bfee09976b3465f/regex-2026.9.3-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:5e674cecb61cb160be392da07fd8a71509ef927f437fbf3215432692ed385151", size = 911828, upload-time = "2026-09-01T00:50:35.72Z" }, - { url = "https://files.pythonhosted.org/packages/4b/8b/9cc6d4123033f7cb82df6cd8ce19eb0fc18a964afe060a03c9b26757c9f3/regex-2026.9.3-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:665207e41bacd435db001099eeab44103197c2c1a729d73ade74688a905ed4ce", size = 801965, upload-time = "2026-09-01T00:50:37.701Z" }, - { url = "https://files.pythonhosted.org/packages/c9/98/39262e91aa87a67c82cbe90a0df4c3d382c7a44811fe80067904085211b4/regex-2026.9.3-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:7a7ddc9a8ca1795166a1ca80364b8ce74187fc210e112d3fb048b711b934f36c", size = 776192, upload-time = "2026-09-01T00:50:39.57Z" }, - { url = "https://files.pythonhosted.org/packages/24/e9/3bb93fe4ee4b6f8ce7ba69b527c4a63cfa3393fc425ab26486041fe441c8/regex-2026.9.3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:e3037d02425863ce9501afbaa04ba967162810004bacde39a53ea9a5b740eb32", size = 785053, upload-time = "2026-09-01T00:50:41.156Z" }, - { url = "https://files.pythonhosted.org/packages/f9/05/31d5bc2553a700c0dfc6b5b6a13c61cdcd1210fde1e304cfa18a33f138b2/regex-2026.9.3-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:3de4eab8c763393b75bbb26f81934ab2cc8794f48f79e90622e3ab7ea57f3d14", size = 860546, upload-time = "2026-09-01T00:50:42.746Z" }, - { url = "https://files.pythonhosted.org/packages/65/a3/2e1e854d80becda0f061093805bbfc037a5849448f46d0a2b71a070d45e2/regex-2026.9.3-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:98620c9c4c22568ad70f57b80527c780b6f8fd26e36507bf8e2273262a228275", size = 765841, upload-time = "2026-09-01T00:50:44.5Z" }, - { url = "https://files.pythonhosted.org/packages/6a/d6/43d02948cedde2e8476ac893ea02755ee5ee1b21c531fda92d80e114f0bc/regex-2026.9.3-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:0b1ba3aaaf5776de473ee16625ac60ac195abb0343afb273575a8201d99be089", size = 852147, upload-time = "2026-09-01T00:50:46.474Z" }, - { url = "https://files.pythonhosted.org/packages/21/ff/adb4e2d08afe8f4c6df004d94604257e1f72af7ba328af7715601585aba4/regex-2026.9.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:56d8659c65166641d8f1b5efccc391c62c8a899eff4d528b981cc62b7b402a4b", size = 789761, upload-time = "2026-09-01T00:50:48.749Z" }, - { url = "https://files.pythonhosted.org/packages/c3/e1/1490d1351758e87f6e702cf2025036bdc7bc59182e2ff5c7bec004b19aed/regex-2026.9.3-cp312-cp312-win32.whl", hash = "sha256:837c1859913798d8bebcd98d4a037e113f8d79e81733009bf590e449769eecb3", size = 267150, upload-time = "2026-09-01T00:50:50.414Z" }, - { url = "https://files.pythonhosted.org/packages/d5/49/4c40cf722d84d60e807a08ef4c3f579216bf97df60c4a1b10be49655d302/regex-2026.9.3-cp312-cp312-win_amd64.whl", hash = "sha256:1ba1dbbb93c5c5629c1861763aec5bfa9f05ad24ef450694130e25029ce7bc36", size = 277773, upload-time = "2026-09-01T00:50:51.963Z" }, - { url = "https://files.pythonhosted.org/packages/aa/af/c48b3b2b4244b4b090554c78d3387e9ae7b859f3dbf7148a27d427e9e5b8/regex-2026.9.3-cp312-cp312-win_arm64.whl", hash = "sha256:d7b3a8a4bbd83ad8b29758f5d24bab10a3f2de87970db36f1e3651c733353136", size = 277122, upload-time = "2026-09-01T00:50:53.778Z" }, -] - -[[package]] -name = "setuptools" -version = "81.0.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/0d/1c/73e719955c59b8e424d015ab450f51c0af856ae46ea2da83eba51cc88de1/setuptools-81.0.0.tar.gz", hash = "sha256:487b53915f52501f0a79ccfd0c02c165ffe06631443a886740b91af4b7a5845a", size = 1198299, upload-time = "2026-02-06T21:10:39.601Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/e1/e3/c164c88b2e5ce7b24d667b9bd83589cf4f3520d97cad01534cd3c4f55fdb/setuptools-81.0.0-py3-none-any.whl", hash = "sha256:fdd925d5c5d9f62e4b74b30d6dd7828ce236fd6ed998a08d81de62ce5a6310d6", size = 1062021, upload-time = "2026-02-06T21:10:37.175Z" }, -] - -[[package]] -name = "sympy" -version = "1.14.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "mpmath" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/83/d3/803453b36afefb7c2bb238361cd4ae6125a569b4db67cd9e79846ba2d68c/sympy-1.14.0.tar.gz", hash = "sha256:d3d3fe8df1e5a0b42f0e7bdf50541697dbe7d23746e894990c030e2b05e72517", size = 7793921, upload-time = "2025-04-27T18:05:01.611Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/a2/09/77d55d46fd61b4a135c444fc97158ef34a095e5681d0a6c10b75bf356191/sympy-1.14.0-py3-none-any.whl", hash = "sha256:e091cc3e99d2141a0ba2847328f5479b05d94a6635cb96148ccb3f34671bd8f5", size = 6299353, upload-time = "2025-04-27T18:04:59.103Z" }, -] - -[[package]] -name = "torch" -version = "2.11.0+cu128" -source = { registry = "https://download.pytorch.org/whl/cu128" } -dependencies = [ - { name = "cuda-bindings", marker = "sys_platform == 'linux'" }, - { name = "cuda-toolkit", extra = ["cublas", "cudart", "cufft", "cufile", "cupti", "curand", "cusolver", "cusparse", "nvjitlink", "nvrtc", "nvtx"], marker = "sys_platform == 'linux'" }, - { name = "filelock" }, - { name = "fsspec" }, - { name = "jinja2" }, - { name = "networkx" }, - { name = "nvidia-cudnn-cu12", marker = "sys_platform == 'linux'" }, - { name = "nvidia-cusparselt-cu12", marker = "sys_platform == 'linux'" }, - { name = "nvidia-nccl-cu12", marker = "sys_platform == 'linux'" }, - { name = "nvidia-nvshmem-cu12", marker = "sys_platform == 'linux'" }, - { name = "setuptools" }, - { name = "sympy" }, - { name = "triton", marker = "sys_platform == 'linux'" }, - { name = "typing-extensions" }, -] -wheels = [ - { url = "https://download-r2.pytorch.org/whl/cu128/torch-2.11.0%2Bcu128-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:9c8f38efee365cb9d334de8a83ce52fc7e5fc9e5a7b0853285efa1b69e00b0f2", upload-time = "2026-04-27T17:41:30Z" }, - { url = "https://download-r2.pytorch.org/whl/cu128/torch-2.11.0%2Bcu128-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:d252cf975fb18c94a85336323ad425f473df56dab35a44b00399bd70c7a3b997", upload-time = "2026-04-27T17:42:06Z" }, - { url = "https://download-r2.pytorch.org/whl/cu128/torch-2.11.0%2Bcu128-cp312-cp312-win_amd64.whl", hash = "sha256:7c78215c3af4f62e63f2b2e360f1722fc719b0853c7ac22666483d9810613a4c", upload-time = "2026-04-27T17:43:49Z" }, -] - -[[package]] -name = "triton" -version = "3.6.0" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/17/5d/08201db32823bdf77a0e2b9039540080b2e5c23a20706ddba942924ebcd6/triton-3.6.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:374f52c11a711fd062b4bfbb201fd9ac0a5febd28a96fb41b4a0f51dde3157f4", size = 176128243, upload-time = "2026-01-20T16:16:07.857Z" }, - { url = "https://files.pythonhosted.org/packages/ab/a8/cdf8b3e4c98132f965f88c2313a4b493266832ad47fb52f23d14d4f86bb5/triton-3.6.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:74caf5e34b66d9f3a429af689c1c7128daba1d8208df60e81106b115c00d6fca", size = 188266850, upload-time = "2026-01-20T16:00:43.041Z" }, -] - -[[package]] -name = "typing-extensions" -version = "4.16.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/f6/cc/6253133b5bb138fc3306cebfbda2c520f545d36b5be2c7255cc528bb45d6/typing_extensions-4.16.0.tar.gz", hash = "sha256:dc983d19a509c94dba722ee6abd33940f7c05a89e243c47e907eb4db6f1a43e5", size = 113555, upload-time = "2026-07-02T08:40:05.92Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/49/d3/b8441a820a491ddfc024b0b0cf0393375b75ea13866d9c66727e54c2fc80/typing_extensions-4.16.0-py3-none-any.whl", hash = "sha256:481caa481374e813c1b176ada14e97f1f67a4539ce9cfeb3f350d78d6370c2e8", size = 45571, upload-time = "2026-07-02T08:40:04.659Z" }, -] - -[[package]] -name = "wcwidth" -version = "0.8.3" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/36/57/ed58088fafdf4c55a0ad6bde846502567645424d7ebf325230b9237f4085/wcwidth-0.8.3.tar.gz", hash = "sha256:d128512515fbf4612e0ff21fd6380399210318b7b54a9af59dff8454cf9730eb", size = 1458450, upload-time = "2026-08-28T18:10:06.875Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/c4/0e/57f6bb3024a597b2e8ec4aee710ffe62ddc95af2e2bb1ee7a7abdc22c68c/wcwidth-0.8.3-py3-none-any.whl", hash = "sha256:d5b73dba6158a595ec9370350e7f2637bcac8d6c5e4fde34f30fcffb6103a5e4", size = 331669, upload-time = "2026-08-28T18:10:04.909Z" }, -] diff --git a/dimos/imitation/policy/abc/test_module.py b/dimos/imitation/policy/abc/test_module.py deleted file mode 100644 index b98df335b3..0000000000 --- a/dimos/imitation/policy/abc/test_module.py +++ /dev/null @@ -1,45 +0,0 @@ -# Copyright 2026 Dimensional Inc. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -from pathlib import Path - -from dimos.experimental.isolated_python.module import contract_rpc_names -from dimos.imitation.policy.abc.module import DualOpenYamAbcPolicy - - -def test_abc_contract_is_importable_without_torch() -> None: - blueprint = DualOpenYamAbcPolicy.blueprint(artifact="checkpoint.pt", task="bottles") - streams = {stream.name for stream in blueprint.blueprints[0].streams} - - assert streams == { - "button_pressed", - "top_image", - "left_wrist_image", - "right_wrist_image", - "coordinator_joint_state", - } - assert contract_rpc_names(DualOpenYamAbcPolicy) == { - "preflight_rollout", - "rollout_status", - "start_rollout", - "stop_rollout", - } - - -def test_abc_contract_resolves_its_own_isolated_project() -> None: - module = DualOpenYamAbcPolicy(artifact="checkpoint.pt", task="bottles") - try: - assert module.runtime_project == Path(__file__).parent / "python" - finally: - module.stop() diff --git a/dimos/imitation/policy/backend.py b/dimos/imitation/policy/backend.py deleted file mode 100644 index d0db762265..0000000000 --- a/dimos/imitation/policy/backend.py +++ /dev/null @@ -1,53 +0,0 @@ -# Copyright 2026 Dimensional Inc. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Small in-process contract implemented by isolated policy backends.""" - -from __future__ import annotations - -from collections.abc import Mapping -from dataclasses import dataclass -from typing import Any, Protocol - -import numpy as np -from numpy.typing import NDArray - -from dimos.imitation.profile import PolicyIOProfile - - -@dataclass(frozen=True) -class PolicyBackendInfo: - """Execution information discovered while loading a policy artifact.""" - - name: str - chunk_length: int - preferred_execution_steps: int - action_lower: NDArray[np.float32] | None = None - action_upper: NDArray[np.float32] | None = None - - -class PolicyBackend(Protocol): - """Backend-specific loading and inference behind the common rollout loop.""" - - def __init__(self, config: Any) -> None: ... - - def load(self, profile: PolicyIOProfile) -> PolicyBackendInfo: ... - - def reset(self) -> None: ... - - def predict( - self, - observations: Mapping[str, NDArray[Any]], - task: str, - ) -> NDArray[np.float32]: ... diff --git a/dimos/imitation/policy/lerobot/README.md b/dimos/imitation/policy/lerobot/README.md index de40e22aef..d53fc8f3e1 100644 --- a/dimos/imitation/policy/lerobot/README.md +++ b/dimos/imitation/policy/lerobot/README.md @@ -1,26 +1,61 @@ -# LeRobot Policy Backend +# LeRobot Policy Module -`OpenYamLeRobotPolicy` is generated from `OPENYAM_QUEST_IO`. Its host process -contains no LeRobot imports; the sibling locked project implements the common -policy backend and runs prediction in the same isolated process as the shared -rollout loop. +`LeRobotPolicyModule` runs trained LeRobot policies in a managed Python-native +subprocess. Its LeRobot, Transformers, Torch, and NumPy versions live in the +sibling `python/` project and do not change the main DimOS environment. + +The host contract subscribes to: + +- `color_image: Image` +- `coordinator_joint_state: JointState` +- `button_pressed: Buttons` + +It submits complete, timestamped action chunks to one named +`JointTrajectoryTask` through the control coordinator. The state and action +vectors use `joint_names` order, including any gripper joint. The policy output +is already postprocessed into each joint's native absolute coordinate; the +runtime does not reinterpret gripper values. ```python -from dimos.imitation.policy.lerobot.module import OpenYamLeRobotPolicy +from dimos.imitation.policy.lerobot.module import LeRobotPolicyModule -policy = OpenYamLeRobotPolicy.blueprint( - instance_name="PolicyRolloutModule", - artifact="outputs/pick/checkpoints/last/pretrained_model", +policy = LeRobotPolicyModule.blueprint( + policy_path="outputs/pick/checkpoints/last/pretrained_model", task="pick up the object", + joint_names=["arm/joint1", "arm/joint2", "arm/gripper"], + trajectory_task_name="policy_rollout", + fps=30.0, + robot_type="my_robot", + image_width=640, + image_height=480, ) ``` -The profile supplies `wrist_image`, `coordinator_joint_state`, feature keys, -image shape, joint order, and 30 Hz rate. The LeRobot adapter validates those -keys and dimensions, loads pre/postprocessors, returns a 2-D action chunk, and -reports checkpoint action bounds to the common safety loop. +The module exposes `preflight_rollout`, `start_rollout`, `stop_rollout`, and +`rollout_status` RPCs. Preflight loads the checkpoint and processors, validates +the control task and fresh live observations, and sends no trajectory. +`start_rollout` refuses to run until preflight passes and rechecks observations +before starting. +The runtime rejects missing or stale observations, missing joints, non-finite +values, incompatible checkpoint features, malformed action chunks, and +trajectories outside the hardware's declared position limits. Pressing the +configured Quest button (A by default) toggles a preflighted rollout. + +The runtime calls LeRobot's `predict_action_chunk()`, postprocesses the entire +chunk, clips every action dimension to the checkpoint's recorded data range, +and executes its first `n_action_steps` at the configured `fps`. The coordinator +still validates the resulting trajectory against hardware limits. Each trajectory +starts with the joint-state observation used for inference, so the coordinator +rejects a stale start if the robot moved in the meantime. Configure `fps` to +match the action frequency used by the training dataset. + +Current limitation: this contract assumes every postprocessed action is an +absolute target in the connected hardware joint's native coordinate. A generic +contract for checkpoints that encode grippers in normalized or device-specific +coordinates remains future work; this runtime does not special-case those +grippers. -Run isolated checks with: +Run isolated runtime checks with: ```bash cd dimos/imitation/policy/lerobot/python diff --git a/dimos/imitation/policy/lerobot/python/dimos_lerobot/runtime.py b/dimos/imitation/policy/lerobot/python/dimos_lerobot/runtime.py index 4f61a78d14..9cd5ce5c48 100644 --- a/dimos/imitation/policy/lerobot/python/dimos_lerobot/runtime.py +++ b/dimos/imitation/policy/lerobot/python/dimos_lerobot/runtime.py @@ -12,13 +12,14 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""LeRobot adapter for the shared DimOS policy rollout runtime.""" +"""Run trained LeRobot policies in an isolated Python environment.""" from __future__ import annotations -from collections.abc import Mapping from contextlib import nullcontext from dataclasses import dataclass +from threading import Condition, Event, RLock, Thread, current_thread +import time from typing import Any from lerobot.configs.policies import PreTrainedConfig @@ -30,15 +31,30 @@ from lerobot.utils.import_utils import register_third_party_plugins import numpy as np from numpy.typing import NDArray +from reactivex.disposable import Disposable import torch -from dimos.imitation.policy.backend import PolicyBackendInfo +from dimos.constants import DEFAULT_THREAD_JOIN_TIMEOUT +from dimos.control.tasks.trajectory_task.trajectory_task import TrajectoryExecutionStatus +from dimos.core.core import rpc from dimos.imitation.policy.lerobot.module import ( - LeRobotPolicyConfig, - OpenYamLeRobotPolicy, + LeRobotPolicyModule, + RolloutStatus, ) -from dimos.imitation.policy.runtime import declare_policy_runtime -from dimos.imitation.profile import ImageSource, PolicyIOProfile +from dimos.msgs.sensor_msgs.Image import Image, ImageFormat +from dimos.msgs.sensor_msgs.JointState import JointState +from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory +from dimos.msgs.trajectory_msgs.TrajectoryPoint import TrajectoryPoint +from dimos.teleop.quest.quest_types import BUTTON_ALIASES, Buttons +from dimos.utils.logging_config import setup_logger + +logger = setup_logger() + +_IMAGE_FEATURE = "observation.images.wrist" +_STATE_FEATURE = "observation.state" +_ACTION_FEATURE = "action" + +RawObservation = dict[str, NDArray[np.uint8] | NDArray[np.float32]] @dataclass(frozen=True) @@ -48,130 +64,472 @@ class _LoadedPolicy: preprocessor: PolicyProcessorPipeline[RobotObservation, RobotObservation] postprocessor: PolicyProcessorPipeline[PolicyAction, PolicyAction] use_amp: bool + chunk_size: int | None + n_action_steps: int + action_lower: NDArray[np.float32] + action_upper: NDArray[np.float32] + + +class LeRobotPolicyRuntime(LeRobotPolicyModule): + """Concrete LeRobot implementation loaded by ``LeRobotPolicyModule``.""" + + _lock: RLock + _observation_changed: Condition + _loaded_policy: _LoadedPolicy | None + _latest_image: tuple[NDArray[np.uint8], float] | None + _latest_joint_state: JointState | None + _stop_event: Event + _thread: Thread | None + _chunks_accepted: int + _last_error: str | None + _active: bool + + def __init__(self, **kwargs: Any) -> None: + super().__init__(**kwargs) + self._lock = RLock() + self._observation_changed = Condition(self._lock) + self._loaded_policy = None + self._latest_image = None + self._latest_joint_state = None + self._stop_event = Event() + self._thread = None + self._chunks_accepted = 0 + self._last_error = None + self._active = False + + @rpc + def start(self) -> None: + super().start() + self.register_disposable(Disposable(self.color_image.subscribe(self._on_color_image))) + self.register_disposable( + Disposable(self.coordinator_joint_state.subscribe(self._on_joint_state)) + ) + self.register_disposable(Disposable(self.button_pressed.subscribe(self._on_button_pressed))) + + @rpc + def stop(self) -> None: + if not self._stop_policy(): + self._cancel_after_stop_timeout() + super().stop() + + @rpc + def preflight_rollout(self) -> RolloutStatus: + """Validate the checkpoint, coordinator, and live observations without moving.""" + with self._lock: + if self._active: + self._last_error = "cannot preflight while a policy rollout is active" + return self._status_locked() + loaded_policy = self._loaded_policy + try: + self._snapshot_observation(time.time()) + except Exception as exc: + self._loaded_policy = None + self._last_error = str(exc) + return self._status_locked() + + try: + tasks = set(self._control.list_tasks()) + if self.config.trajectory_task_name not in tasks: + raise RuntimeError( + "ControlCoordinator is missing configured rollout task " + f"{self.config.trajectory_task_name!r}" + ) + if loaded_policy is None: + loaded_policy = self._load_policy() + logger.info( + "Loaded LeRobot policy during preflight", + path=self.config.policy_path, + runtime_fps=self.config.fps, + chunk_size=loaded_policy.chunk_size, + n_action_steps=loaded_policy.n_action_steps, + ) + with self._lock: + self._loaded_policy = loaded_policy + self._snapshot_observation(time.time()) + self._last_error = None + return self._status_locked() + except Exception as exc: + with self._lock: + self._loaded_policy = None + self._last_error = str(exc) + return self._status_locked() + @rpc + def start_rollout(self) -> RolloutStatus: + with self._lock: + if self._thread is not None and self._thread.is_alive(): + self._last_error = "a policy rollout is already active" + return self._status_locked() + if self._loaded_policy is None: + self._last_error = "policy preflight has not passed" + return self._status_locked() + try: + self._snapshot_observation(time.time()) + except RuntimeError as exc: + self._last_error = str(exc) + return self._status_locked() -class LeRobotBackend: - """Translate the generic profile-keyed arrays to a LeRobot checkpoint.""" + self._stop_event.clear() + self._chunks_accepted = 0 + self._last_error = None + self._active = True + self._thread = Thread( + target=self._run_rollout, + name="lerobot-policy-rollout", + daemon=True, + ) + self._thread.start() + return self._status_locked() + + @rpc + def stop_rollout(self) -> RolloutStatus: + if not self._stop_policy(): + self._cancel_after_stop_timeout() + return self.rollout_status() + + @rpc + def rollout_status(self) -> RolloutStatus: + with self._lock: + return self._status_locked() + + def _status_locked(self) -> RolloutStatus: + try: + self._snapshot_observation(time.time()) + observations_ready = True + except RuntimeError: + observations_ready = False + return { + "active": self._active, + "policy_path": self.config.policy_path, + "task": self.config.task, + "device": self.config.device, + "policy_ready": self._loaded_policy is not None, + "observations_ready": observations_ready, + "chunks_accepted": self._chunks_accepted, + "last_error": self._last_error, + } + + def _on_color_image(self, image: Image) -> None: + if image.format != ImageFormat.RGB or image.data.dtype != np.uint8: + logger.warning("Ignoring non-uint8 RGB policy image", image=str(image)) + return + expected_shape = (self.config.image_height, self.config.image_width, 3) + if image.data.shape != expected_shape: + logger.warning( + "Ignoring policy image with unexpected shape", + shape=image.data.shape, + expected=expected_shape, + ) + return + with self._lock: + self._latest_image = (np.ascontiguousarray(image.data), image.ts) - def __init__(self, config: LeRobotPolicyConfig) -> None: - self.config = config - self._loaded: _LoadedPolicy | None = None - self._profile: PolicyIOProfile | None = None + def _on_joint_state(self, state: JointState) -> None: + with self._observation_changed: + self._latest_joint_state = JointState(state) + self._observation_changed.notify_all() - def load(self, profile: PolicyIOProfile) -> PolicyBackendInfo: + def _on_button_pressed(self, buttons: Buttons) -> None: + button = BUTTON_ALIASES.get(self.config.rollout_button, self.config.rollout_button) + if not bool(getattr(buttons, button)): + return + with self._lock: + active = self._active + if active: + self.stop_rollout() + else: + self.start_rollout() + + def _snapshot_observation( + self, now: float + ) -> tuple[NDArray[np.uint8], NDArray[np.float32], float]: + if self._latest_image is None: + raise RuntimeError("no camera image has been received") + if self._latest_joint_state is None: + raise RuntimeError("no coordinator joint state has been received") + + image, image_ts = self._latest_image + state = self._latest_joint_state + max_age = self.config.max_observation_age_s + if now - image_ts > max_age: + raise RuntimeError(f"camera image is stale by {now - image_ts:.2f}s") + if now - state.ts > max_age: + raise RuntimeError(f"joint state is stale by {now - state.ts:.2f}s") + + positions = dict(zip(state.name, state.position, strict=False)) + missing = [name for name in self.config.joint_names if name not in positions] + if missing: + raise RuntimeError(f"joint state is missing configured joints: {missing}") + vector = np.asarray( + [positions[name] for name in self.config.joint_names], + dtype=np.float32, + ) + if not np.all(np.isfinite(vector)): + raise RuntimeError("joint state contains non-finite positions") + return image.copy(), vector, state.ts + + def _load_policy(self) -> _LoadedPolicy: register_third_party_plugins() - policy_config = PreTrainedConfig.from_pretrained(self.config.artifact) + policy_config = PreTrainedConfig.from_pretrained(self.config.policy_path) if self.config.device is not None: policy_config.device = self.config.device if policy_config.device is None: raise RuntimeError("LeRobot did not resolve an inference device") - _validate_features(policy_config, profile) + self._validate_features(policy_config) device = torch.device(policy_config.device) if device.type == "cuda" and not torch.cuda.is_available(): raise RuntimeError( f"Policy requested device {policy_config.device!r}, but CUDA is not available" ) + policy_class = get_policy_class(policy_config.type) - policy = policy_class.from_pretrained(self.config.artifact, config=policy_config) + loaded_policy = policy_class.from_pretrained(self.config.policy_path, config=policy_config) preprocessor, postprocessor = make_pre_post_processors( policy_cfg=policy_config, - pretrained_path=self.config.artifact, + pretrained_path=self.config.policy_path, preprocessor_overrides={"device_processor": {"device": str(device)}}, ) - width = len(profile.action.demonstration.joints) - lower, upper = _checkpoint_action_bounds(postprocessor, width) - n_action_steps = _positive_int_attribute(policy_config, "n_action_steps") - chunk_length = _optional_int_attribute(policy_config, "chunk_size") or n_action_steps - self._loaded = _LoadedPolicy( - policy=policy, + action_lower, action_upper = _checkpoint_action_bounds( + postprocessor, + len(self.config.joint_names), + ) + return _LoadedPolicy( + policy=loaded_policy, device=device, preprocessor=preprocessor, postprocessor=postprocessor, use_amp=bool(policy_config.use_amp), + chunk_size=_optional_int_attribute(policy_config, "chunk_size"), + n_action_steps=_positive_int_attribute(policy_config, "n_action_steps"), + action_lower=action_lower, + action_upper=action_upper, ) - self._profile = profile - return PolicyBackendInfo( - name="lerobot", - chunk_length=chunk_length, - preferred_execution_steps=n_action_steps, - action_lower=lower, - action_upper=upper, - ) - def reset(self) -> None: - loaded = self._require_loaded() - _reset(loaded.policy) - _reset(loaded.preprocessor) - _reset(loaded.postprocessor) + def _validate_features(self, policy_config: PreTrainedConfig) -> None: + inputs = policy_config.input_features or {} + outputs = policy_config.output_features or {} + missing = {_IMAGE_FEATURE, _STATE_FEATURE} - set(inputs) + if missing: + raise ValueError( + "Policy is incompatible with the DimOS single-camera runtime; " + f"missing input features: {sorted(missing)}" + ) + if _ACTION_FEATURE not in outputs: + raise ValueError(f"Policy has no {_ACTION_FEATURE!r} output feature") + if getattr(policy_config, "temporal_ensemble_coeff", None) is not None: + raise ValueError("Policies using temporal ensembling are not supported") + + state_shape = tuple(inputs[_STATE_FEATURE].shape) + image_shape = tuple(inputs[_IMAGE_FEATURE].shape) + action_shape = tuple(outputs[_ACTION_FEATURE].shape) + joint_count = len(self.config.joint_names) + expected_image_shape = (3, self.config.image_height, self.config.image_width) + if image_shape != expected_image_shape: + raise ValueError( + f"Policy image shape {image_shape} does not match {expected_image_shape}" + ) + if not state_shape or state_shape[0] != joint_count: + raise ValueError( + f"Policy state dimension {state_shape} does not match {joint_count} configured joints" + ) + if not action_shape or action_shape[0] != joint_count: + raise ValueError( + f"Policy action dimension {action_shape} does not match {joint_count} configured joints" + ) - def predict( + def _predict( self, - observations: Mapping[str, NDArray[Any]], + loaded_policy: _LoadedPolicy, + image: NDArray[np.uint8], + state: NDArray[np.float32], + *, task: str, ) -> NDArray[np.float32]: - loaded = self._require_loaded() - assert self._profile is not None + observation: RawObservation = { + _IMAGE_FEATURE: image, + _STATE_FEATURE: state, + } with ( torch.inference_mode(), torch.autocast(device_type="cuda") - if loaded.device.type == "cuda" and loaded.use_amp + if loaded_policy.device.type == "cuda" and loaded_policy.use_amp else nullcontext(), ): prepared = prepare_observation_for_inference( - dict(observations), - loaded.device, + observation, + loaded_policy.device, task=task, - robot_type=self._profile.robot_type, + robot_type=self.config.robot_type, ) - prepared = loaded.preprocessor(prepared) - predict = getattr(loaded.policy, "predict_action_chunk", None) + prepared = loaded_policy.preprocessor(prepared) + predict = getattr(loaded_policy.policy, "predict_action_chunk", None) if not callable(predict): raise TypeError("Policy does not provide predict_action_chunk()") - action_chunk = loaded.postprocessor(predict(prepared)) - result = np.asarray(action_chunk.to("cpu").numpy(), dtype=np.float32) - if result.ndim != 3 or result.shape[0] != 1: - raise RuntimeError(f"LeRobot returned invalid batched action shape {result.shape}") - return np.asarray(result[0], dtype=np.float32) - - def _require_loaded(self) -> _LoadedPolicy: - if self._loaded is None: - raise RuntimeError("LeRobot backend is not loaded") - return self._loaded - - -LeRobotPolicyRuntime = declare_policy_runtime( - "LeRobotPolicyRuntime", - __name__, - OpenYamLeRobotPolicy, - LeRobotBackend, -) + action_chunk = loaded_policy.postprocessor(predict(prepared)) + return np.asarray(action_chunk.to("cpu").numpy(), dtype=np.float32) + + def _run_rollout(self) -> None: + loaded_policy: _LoadedPolicy | None = None + try: + with self._lock: + loaded_policy = self._loaded_policy + if loaded_policy is None: + raise RuntimeError("policy preflight has not passed") + if self._stop_event.is_set(): + return + self._reset_policy(loaded_policy) + while not self._stop_event.is_set(): + with self._lock: + image, state, state_ts = self._snapshot_observation(time.time()) + action_chunk = self._predict(loaded_policy, image, state, task=self.config.task) + expected_width = len(self.config.joint_names) + if action_chunk.ndim != 3 or action_chunk.shape[0] != 1: + raise RuntimeError( + f"policy returned action chunk shape {action_chunk.shape}, expected " + f"(1, steps, {expected_width})" + ) + if action_chunk.shape[2] != expected_width: + raise RuntimeError( + f"policy returned action width {action_chunk.shape[2]}, expected {expected_width}" + ) + if action_chunk.shape[1] < loaded_policy.n_action_steps: + raise RuntimeError( + f"policy returned {action_chunk.shape[1]} action steps, but n_action_steps " + f"is {loaded_policy.n_action_steps}" + ) + actions = action_chunk[0, : loaded_policy.n_action_steps] + if not np.all(np.isfinite(actions)): + raise RuntimeError("policy returned non-finite joint targets") + bounded_actions = np.clip( + actions, + loaded_policy.action_lower, + loaded_policy.action_upper, + ) + clipped = np.any(actions != bounded_actions, axis=0) + if np.any(clipped): + logger.warning( + "Clipped policy actions to checkpoint range", + joints=[ + name + for name, was_clipped in zip( + self.config.joint_names, clipped, strict=True + ) + if was_clipped + ], + ) + actions = bounded_actions + if self._stop_event.is_set(): + break + result = self._control.execute_trajectory( + self._trajectory(state, actions), + task_name=self.config.trajectory_task_name, + ) + if result.status is TrajectoryExecutionStatus.START_STATE_MISMATCH: + self._wait_for_newer_joint_state(state_ts) + continue + if result.status is not TrajectoryExecutionStatus.ACCEPTED: + raise RuntimeError( + result.message or f"trajectory rejected: {result.status.name}" + ) + with self._lock: + self._chunks_accepted += 1 + self._stop_event.wait(loaded_policy.n_action_steps / self.config.fps) + except Exception as exc: + with self._lock: + self._last_error = str(exc) + logger.exception("LeRobot policy execution stopped", error=str(exc)) + finally: + self._stop_event.set() + cancellation_error = self._cancel_trajectory() + if loaded_policy is not None: + self._reset_policy(loaded_policy) + with self._lock: + if cancellation_error is not None: + self._last_error = ( + f"{self._last_error}; {cancellation_error}" + if self._last_error is not None + else cancellation_error + ) + self._active = False -def _validate_features(policy_config: PreTrainedConfig, profile: PolicyIOProfile) -> None: - inputs = policy_config.input_features or {} - outputs = policy_config.output_features or {} - missing = set(profile.observations) - set(inputs) - if missing: - raise ValueError(f"LeRobot checkpoint is missing input features: {sorted(missing)}") - if profile.action.key not in outputs: - raise ValueError(f"LeRobot checkpoint has no {profile.action.key!r} output feature") - if getattr(policy_config, "temporal_ensemble_coeff", None) is not None: - raise ValueError("Policies using temporal ensembling are not supported") - - for key, source in profile.observations.items(): - actual = tuple(inputs[key].shape) - expected = ( - (source.shape[2], source.shape[0], source.shape[1]) - if isinstance(source, ImageSource) - else (len(source.joints),) + @staticmethod + def _reset_policy(loaded_policy: _LoadedPolicy) -> None: + _reset(loaded_policy.policy) + _reset(loaded_policy.preprocessor) + _reset(loaded_policy.postprocessor) + + def _stop_policy(self) -> bool: + with self._lock: + thread = self._thread + self._stop_event.set() + self._observation_changed.notify_all() + if thread is not None and thread is not current_thread(): + thread.join(timeout=DEFAULT_THREAD_JOIN_TIMEOUT) + return thread is None or not thread.is_alive() + + def _cancel_after_stop_timeout(self) -> None: + timeout_error = f"policy rollout did not stop within {DEFAULT_THREAD_JOIN_TIMEOUT} seconds" + cancellation_error = self._cancel_trajectory() + with self._lock: + self._last_error = ( + f"{timeout_error}; {cancellation_error}" + if cancellation_error is not None + else timeout_error + ) + + def _trajectory( + self, + state: NDArray[np.float32], + actions: NDArray[np.float32], + ) -> JointTrajectory: + zeros = [0.0] * len(self.config.joint_names) + points = [ + TrajectoryPoint( + positions=[float(value) for value in state], + velocities=zeros, + time_from_start=0.0, + ) + ] + points.extend( + TrajectoryPoint( + positions=[float(value) for value in action], + velocities=zeros, + time_from_start=(index + 1) / self.config.fps, + ) + for index, action in enumerate(actions) ) - if actual != expected: - raise ValueError(f"LeRobot feature {key!r} shape {actual} does not match {expected}") - action_shape = tuple(outputs[profile.action.key].shape) - expected_action = (len(profile.action.demonstration.joints),) - if action_shape != expected_action: - raise ValueError(f"LeRobot action shape {action_shape} does not match {expected_action}") + return JointTrajectory(joint_names=list(self.config.joint_names), points=points) + + def _wait_for_newer_joint_state(self, previous_ts: float) -> None: + with self._observation_changed: + self._observation_changed.wait_for( + lambda: self._stop_event.is_set() + or ( + self._latest_joint_state is not None + and self._latest_joint_state.ts > previous_ts + ) + ) + + def _cancel_trajectory(self) -> str | None: + try: + result = self._control.cancel_trajectory(task_name=self.config.trajectory_task_name) + except Exception as exc: + logger.exception( + "Failed to cancel policy trajectory", + task_name=self.config.trajectory_task_name, + ) + return f"Failed to cancel policy trajectory: {exc}" + if result.safe: + return None + message = result.message or "Policy trajectory cancellation was uncertain" + logger.error( + "Policy trajectory cancellation was uncertain", + error=message, + task_name=self.config.trajectory_task_name, + ) + return message def _checkpoint_action_bounds( @@ -188,11 +546,19 @@ def _checkpoint_action_bounds( break if lower_tensor is None or upper_tensor is None: raise ValueError("Policy postprocessor has no action min/max statistics") + lower = np.asarray(lower_tensor.detach().cpu().numpy(), dtype=np.float32) upper = np.asarray(upper_tensor.detach().cpu().numpy(), dtype=np.float32) - shape = (expected_width,) - if lower.shape != shape or upper.shape != shape: - raise ValueError(f"Policy action range shape must be {shape}") + expected_shape = (expected_width,) + if lower.shape != expected_shape or upper.shape != expected_shape: + raise ValueError( + "Policy action range shape does not match configured joints: " + f"min={lower.shape}, max={upper.shape}, expected={expected_shape}" + ) + if not np.all(np.isfinite(lower)) or not np.all(np.isfinite(upper)): + raise ValueError("Policy action range contains non-finite values") + if np.any(lower > upper): + raise ValueError("Policy action range has min greater than max") return lower, upper diff --git a/dimos/imitation/policy/lerobot/python/dimos_lerobot/runtime_tests.py b/dimos/imitation/policy/lerobot/python/dimos_lerobot/runtime_tests.py index 406aebeb72..c427b41fd6 100644 --- a/dimos/imitation/policy/lerobot/python/dimos_lerobot/runtime_tests.py +++ b/dimos/imitation/policy/lerobot/python/dimos_lerobot/runtime_tests.py @@ -12,19 +12,36 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Any +"""Behavior tests for the isolated LeRobot runtime.""" -from dimos_lerobot.runtime import ( - LeRobotPolicyRuntime, - _checkpoint_action_bounds, - _validate_features, -) +from collections.abc import Callable, Iterator +from threading import Event, Thread +import time +from typing import Any, Protocol + +from dimos_lerobot import runtime as policy_runtime +from dimos_lerobot.runtime import LeRobotPolicyRuntime +from lerobot.configs.policies import PreTrainedConfig +import numpy as np +from numpy.typing import NDArray import pytest +import pytest_mock import torch +from torch import Tensor + +from dimos.control.tasks.trajectory_task.trajectory_task import ( + TrajectoryCancellationResult, + TrajectoryCancellationStatus, + TrajectoryExecutionResult, + TrajectoryExecutionStatus, +) +from dimos.msgs.sensor_msgs.Image import Image, ImageFormat +from dimos.msgs.sensor_msgs.JointState import JointState +from dimos.protocol.rpc.pubsubrpc import LCMRPC +from dimos.teleop.quest.quest_types import Buttons +from dimos.utils.testing.waiting import wait_until -from dimos.experimental.isolated_python.bootstrap import validate_runtime -from dimos.imitation.policy.lerobot.module import OpenYamLeRobotPolicy -from dimos.robot.manipulators.openyam.learning import OPENYAM_QUEST_IO +JOINTS = [f"test_arm/joint{i}" for i in range(1, 5)] class FakeFeature: @@ -32,45 +49,611 @@ def __init__(self, shape: tuple[int, ...]) -> None: self.shape = shape -class FakeConfig: - temporal_ensemble_coeff = None - input_features = { - "observation.images.wrist": FakeFeature((3, 480, 640)), - "observation.state": FakeFeature((7,)), - } - output_features = {"action": FakeFeature((7,))} +class FakeUpstreamConfig: + def __init__(self, joint_count: int, n_action_steps: int) -> None: + self.type = "fake_policy" + self.device: str | None = "cpu" + self.use_amp = False + self.chunk_size = 3 + self.n_action_steps: int | None = n_action_steps + self.temporal_ensemble_coeff: float | None = None + self.input_features = { + "observation.images.wrist": FakeFeature((3, 4, 5)), + "observation.state": FakeFeature((joint_count,)), + } + self.output_features = {"action": FakeFeature((joint_count,))} + + +class FakePipeline: + def __init__(self, steps: list[object] | None = None) -> None: + self.calls: list[object] = [] + self.reset_count = 0 + self.steps = steps or [] + + def __call__(self, value: object) -> object: + self.calls.append(value) + return value + def reset(self) -> None: + self.reset_count += 1 -class FakeStats: - def state_dict(self) -> dict[str, torch.Tensor]: - return { - "action.min": torch.zeros(7), - "action.max": torch.ones(7), + +class FakeActionStats: + def __init__(self, lower: list[float], upper: list[float]) -> None: + self._state = { + "action.min": torch.tensor(lower), + "action.max": torch.tensor(upper), } + def state_dict(self) -> dict[str, Tensor]: + return self._state -class FakePipeline: - steps: list[Any] = [FakeStats()] +class FakePolicy: + def __init__( + self, + action_chunk: NDArray[np.float32], + n_action_steps: int = 2, + *, + action_lower: list[float] | None = None, + action_upper: list[float] | None = None, + ) -> None: + self.action_chunk = torch.from_numpy(action_chunk).unsqueeze(0) + self.called = Event() + self.reset_count = 0 + self.batch: dict[str, object] | None = None + self.upstream_config = FakeUpstreamConfig(len(JOINTS), n_action_steps) + self.preprocessor = FakePipeline() + self.postprocessor = FakePipeline( + [ + FakeActionStats( + action_lower or [-100.0] * len(JOINTS), + action_upper or [100.0] * len(JOINTS), + ) + ] + ) + self.config_load_count = 0 + + def reset(self) -> None: + self.reset_count += 1 + + def predict_action_chunk(self, batch: dict[str, object]) -> Tensor: + self.batch = dict(batch) + self.called.set() + return self.action_chunk + + +class RuntimeFactory(Protocol): + def __call__( + self, + policy: FakePolicy, + *, + device: str | None = None, + ) -> tuple[LeRobotPolicyRuntime, Any]: ... + + +@pytest.fixture +def make_runtime(mocker: pytest_mock.MockerFixture) -> Iterator[RuntimeFactory]: + mocker.patch("dimos.core.module.get_loop", return_value=(mocker.MagicMock(), None)) + mocker.patch.object(LCMRPC, "__init__", return_value=None) + mocker.patch.object(LCMRPC, "serve_module_rpc", return_value=None) + mocker.patch.object(LCMRPC, "start", return_value=None) + mocker.patch.object(LCMRPC, "stop", return_value=None) + built: list[LeRobotPolicyRuntime] = [] + + def _make( + policy: FakePolicy, + *, + device: str | None = None, + ) -> tuple[LeRobotPolicyRuntime, Any]: + def load_config(_path: str) -> FakeUpstreamConfig: + policy.config_load_count += 1 + return policy.upstream_config + + policy_class = mocker.MagicMock() + policy_class.from_pretrained.return_value = policy + mocker.patch.object(PreTrainedConfig, "from_pretrained", side_effect=load_config) + mocker.patch.object(policy_runtime, "get_policy_class", return_value=policy_class) + mocker.patch.object( + policy_runtime, + "make_pre_post_processors", + return_value=(policy.preprocessor, policy.postprocessor), + ) + + def prepare_observation( + observation: dict[str, NDArray[np.uint8] | NDArray[np.float32]], + _device: torch.device, + *, + task: str, + robot_type: str, + ) -> dict[str, object]: + return { + **{ + name: torch.from_numpy(value).unsqueeze(0) + for name, value in observation.items() + }, + "task": task, + "robot_type": robot_type, + } + + mocker.patch.object( + policy_runtime, + "prepare_observation_for_inference", + side_effect=prepare_observation, + ) + mocker.patch.object(policy_runtime, "register_third_party_plugins") + + module = LeRobotPolicyRuntime( + _isolated_python_runtime=True, + policy_path="checkpoint/default", + task="pick up the test object", + device=device, + joint_names=JOINTS, + fps=50.0, + robot_type="test_arm", + image_width=5, + image_height=4, + ) + control = mocker.MagicMock() + control.execute_trajectory.return_value = TrajectoryExecutionResult( + TrajectoryExecutionStatus.ACCEPTED + ) + control.cancel_trajectory.return_value = TrajectoryCancellationResult( + TrajectoryCancellationStatus.ALREADY_STOPPED + ) + control.list_tasks.return_value = ["policy_rollout"] + mocker.patch.object(module, "_control", control, create=True) + built.append(module) + return module, control + + yield _make + for module in built: + module.stop() + + +def _action_chunk(steps: int = 3) -> NDArray[np.float32]: + return np.arange(steps * len(JOINTS), dtype=np.float32).reshape( + steps, len(JOINTS) + ) / np.float32(10) + + +def _provide_observation( + module: LeRobotPolicyRuntime, + *, + positions: list[float] | None = None, + ts: float | None = None, +) -> tuple[NDArray[np.uint8], list[float], float]: + rgb = np.zeros((4, 5, 3), dtype=np.uint8) + rgb[..., 0] = 10 + rgb[..., 1] = 20 + rgb[..., 2] = 30 + values = positions or [float(i) / 10 for i in range(len(JOINTS))] + timestamp = time.time() if ts is None else ts + module._on_color_image(Image(data=rgb, format=ImageFormat.RGB, ts=timestamp)) + module._on_joint_state(JointState(ts=timestamp, name=JOINTS, position=values)) + return rgb, values, timestamp + + +def _preflight(module: LeRobotPolicyRuntime) -> None: + status = module.preflight_rollout() + assert status["policy_ready"] is True + assert status["observations_ready"] is True + assert status["last_error"] is None + + +def test_policy_predicts_and_executes_one_native_joint_chunk(make_runtime: RuntimeFactory) -> None: + actions = _action_chunk() + policy = FakePolicy(actions, n_action_steps=2) + module, control = make_runtime(policy) + rgb, positions, _ts = _provide_observation(module) + + _preflight(module) + assert module.start_rollout()["active"] is True + wait_until(lambda: control.execute_trajectory.call_count >= 1, timeout=1.0) + module.stop_rollout() + + call = control.execute_trajectory.call_args_list[0] + trajectory = call.args[0] + assert call.kwargs == {"task_name": "policy_rollout"} + assert trajectory.joint_names == JOINTS + assert [point.time_from_start for point in trajectory.points] == [0.0, 0.02, 0.04] + np.testing.assert_allclose(trajectory.points[0].positions, positions) + np.testing.assert_allclose(trajectory.points[1].positions, actions[0]) + np.testing.assert_allclose(trajectory.points[2].positions, actions[1]) + assert policy.batch is not None + assert policy.batch["task"] == "pick up the test object" + image = policy.batch["observation.images.wrist"] + assert isinstance(image, Tensor) + np.testing.assert_array_equal(image.squeeze(0).numpy(), rgb) + assert policy.postprocessor.calls + assert module.rollout_status()["chunks_accepted"] >= 1 + + +def test_policy_actions_are_clipped_to_checkpoint_range(make_runtime: RuntimeFactory) -> None: + actions = np.zeros((3, len(JOINTS)), dtype=np.float32) + actions[:, -1] = 1.016 + policy = FakePolicy( + actions, + n_action_steps=2, + action_lower=[-10.0, -10.0, -10.0, 0.0], + action_upper=[10.0, 10.0, 10.0, 1.0], + ) + module, control = make_runtime(policy) + _provide_observation(module) + + _preflight(module) + module.start_rollout() + wait_until(lambda: control.execute_trajectory.call_count >= 1, timeout=1.0) + module.stop_rollout() + + trajectory = control.execute_trajectory.call_args_list[0].args[0] + assert [point.positions[-1] for point in trajectory.points[1:]] == [1.0, 1.0] + + +def test_next_chunk_uses_latest_joint_observation(make_runtime: RuntimeFactory) -> None: + policy = FakePolicy(_action_chunk(), n_action_steps=1) + module, control = make_runtime(policy) + first_submitted = Event() + release_first = Event() + + def execute_trajectory(*_args: object, **_kwargs: object) -> TrajectoryExecutionResult: + if not first_submitted.is_set(): + first_submitted.set() + assert release_first.wait(timeout=1.0) + return TrajectoryExecutionResult(TrajectoryExecutionStatus.ACCEPTED) + + control.execute_trajectory.side_effect = execute_trajectory + _provide_observation(module, positions=[0.0] * len(JOINTS)) + _preflight(module) + module.start_rollout() + assert first_submitted.wait(timeout=1.0) + + latest = [0.4, 0.3, 0.2, 0.1] + _provide_observation(module, positions=latest) + release_first.set() + wait_until(lambda: control.execute_trajectory.call_count >= 2, timeout=1.0) + module.stop_rollout() + + second = control.execute_trajectory.call_args_list[1].args[0] + np.testing.assert_allclose(second.points[0].positions, latest) + + +def test_start_mismatch_waits_for_new_joint_state_before_retry( + make_runtime: RuntimeFactory, +) -> None: + policy = FakePolicy(_action_chunk(), n_action_steps=1) + module, control = make_runtime(policy) + _provide_observation(module) + + def execute_trajectory(*_args: object, **_kwargs: object) -> TrajectoryExecutionResult: + status = ( + TrajectoryExecutionStatus.START_STATE_MISMATCH + if control.execute_trajectory.call_count == 1 + else TrajectoryExecutionStatus.ACCEPTED + ) + return TrajectoryExecutionResult(status) + + control.execute_trajectory.side_effect = execute_trajectory + _preflight(module) + module.start_rollout() + wait_until(lambda: control.execute_trajectory.call_count == 1, timeout=1.0) + + _provide_observation(module, positions=[0.2] * len(JOINTS), ts=time.time() + 0.01) + wait_until(lambda: control.execute_trajectory.call_count >= 2, timeout=1.0) + module.stop_rollout() + + assert module.rollout_status()["last_error"] is None + + +def test_a_press_stops_worker_before_cancelling_its_trajectory( + make_runtime: RuntimeFactory, +) -> None: + policy = FakePolicy(_action_chunk(), n_action_steps=2) + module, control = make_runtime(policy) + _provide_observation(module) + _preflight(module) + execute_started = Event() + release_execute = Event() + stop_finished = Event() + + def execute_trajectory(*_args: object, **_kwargs: object) -> TrajectoryExecutionResult: + execute_started.set() + assert release_execute.wait(timeout=1.0) + return TrajectoryExecutionResult(TrajectoryExecutionStatus.ACCEPTED) + + control.execute_trajectory.side_effect = execute_trajectory + pressed = Buttons() + pressed.right_primary = True + + def stop_from_button() -> None: + module._on_button_pressed(pressed) + stop_finished.set() -def test_generated_runtime_implements_the_host_contract() -> None: - validate_runtime(OpenYamLeRobotPolicy, LeRobotPolicyRuntime) + module._on_button_pressed(pressed) + assert execute_started.wait(timeout=1.0) + stop_thread = Thread(target=stop_from_button) + stop_thread.start() + try: + assert module._stop_event.wait(timeout=1.0) + assert module.rollout_status()["active"] is True + control.cancel_trajectory.assert_not_called() + finally: + release_execute.set() + stop_thread.join(timeout=1.0) + assert stop_finished.is_set() + assert control.execute_trajectory.call_count == 1 + control.cancel_trajectory.assert_called_with(task_name="policy_rollout") -def test_lerobot_feature_validation_uses_profile_keys_and_shapes() -> None: - _validate_features(FakeConfig(), OPENYAM_QUEST_IO) # type: ignore[arg-type] +def test_uncertain_cancellation_is_reported(make_runtime: RuntimeFactory) -> None: + policy = FakePolicy(_action_chunk(), n_action_steps=1) + module, control = make_runtime(policy) + _provide_observation(module) + _preflight(module) + control.cancel_trajectory.return_value = TrajectoryCancellationResult( + TrajectoryCancellationStatus.UNCERTAIN, + "coordinator did not confirm cancellation", + ) + + module.start_rollout() + wait_until(lambda: control.execute_trajectory.call_count >= 1, timeout=1.0) + status = module.stop_rollout() + + assert status["active"] is False + assert status["last_error"] == "coordinator did not confirm cancellation" + + +@pytest.mark.parametrize( + ("actions", "message"), + [ + (np.zeros((3, len(JOINTS) - 1), dtype=np.float32), "action width"), + (np.full((3, len(JOINTS)), np.nan, dtype=np.float32), "non-finite joint targets"), + ], +) +def test_invalid_action_chunk_cancels_and_latches_rollout_off( + make_runtime: RuntimeFactory, + actions: NDArray[np.float32], + message: str, +) -> None: + policy = FakePolicy(actions) + module, control = make_runtime(policy) + _provide_observation(module) + + _preflight(module) + module.start_rollout() + + wait_until(lambda: module.rollout_status()["active"] is False, timeout=1.0) + assert message in (module.rollout_status()["last_error"] or "") + control.cancel_trajectory.assert_called_with(task_name="policy_rollout") + + +def test_trajectory_rejection_cancels_and_latches_rollout_off( + make_runtime: RuntimeFactory, +) -> None: + policy = FakePolicy(_action_chunk()) + module, control = make_runtime(policy) + _provide_observation(module) + control.execute_trajectory.return_value = TrajectoryExecutionResult( + TrajectoryExecutionStatus.POSITION_LIMIT_VIOLATION, + "outside hardware limits", + ) + + _preflight(module) + module.start_rollout() + + wait_until(lambda: module.rollout_status()["active"] is False, timeout=1.0) + assert module.rollout_status()["last_error"] == "outside hardware limits" + control.cancel_trajectory.assert_called_with(task_name="policy_rollout") + + +@pytest.mark.parametrize( + ("configure", "message"), + [ + (lambda config: setattr(config, "n_action_steps", None), "positive int"), + (lambda config: setattr(config, "n_action_steps", 0), "positive int"), + ( + lambda config: setattr(config, "temporal_ensemble_coeff", 0.01), + "temporal ensembling", + ), + ( + lambda config: setattr( + config.input_features["observation.images.wrist"], + "shape", + (3, 8, 8), + ), + "Policy image shape", + ), + ], +) +def test_incompatible_chunk_contract_is_rejected( + make_runtime: RuntimeFactory, + configure: Callable[[FakeUpstreamConfig], None], + message: str, +) -> None: + policy = FakePolicy(_action_chunk()) + configure(policy.upstream_config) + module, control = make_runtime(policy) + _provide_observation(module) + + status = module.preflight_rollout() + + assert status["active"] is False + assert status["policy_ready"] is False + assert message in (status["last_error"] or "") + control.execute_trajectory.assert_not_called() + + +def test_policy_refuses_to_load_without_live_observations(make_runtime: RuntimeFactory) -> None: + policy = FakePolicy(_action_chunk()) + module, control = make_runtime(policy) + + result = module.preflight_rollout() + + assert result["active"] is False + assert "no camera image" in (result["last_error"] or "") + assert policy.config_load_count == 0 + control.execute_trajectory.assert_not_called() + + +def test_policy_refuses_stale_observations(make_runtime: RuntimeFactory) -> None: + policy = FakePolicy(_action_chunk()) + module, control = make_runtime(policy) + _provide_observation(module, ts=time.time() - module.config.max_observation_age_s - 1.0) + + result = module.preflight_rollout() + + assert result["active"] is False + assert "camera image is stale" in (result["last_error"] or "") + assert policy.config_load_count == 0 + control.execute_trajectory.assert_not_called() + + +def test_checkpoint_loads_on_demand_and_is_cached(make_runtime: RuntimeFactory) -> None: + policy = FakePolicy(_action_chunk()) + module, control = make_runtime(policy) + _provide_observation(module) + + _preflight(module) + module.start_rollout() + wait_until(lambda: control.execute_trajectory.call_count >= 1, timeout=1.0) + module.stop_rollout() + control.execute_trajectory.reset_mock() + module.start_rollout() + wait_until(lambda: control.execute_trajectory.call_count >= 1, timeout=1.0) + module.stop_rollout() + + assert policy.config_load_count == 1 + + +def test_start_requires_a_successful_preflight(make_runtime: RuntimeFactory) -> None: + policy = FakePolicy(_action_chunk()) + module, control = make_runtime(policy) + _provide_observation(module) + + status = module.start_rollout() + + assert status["active"] is False + assert status["last_error"] == "policy preflight has not passed" + control.execute_trajectory.assert_not_called() + + +def test_preflight_never_sends_a_trajectory(make_runtime: RuntimeFactory) -> None: + policy = FakePolicy(_action_chunk()) + module, control = make_runtime(policy) + _provide_observation(module) + + status = module.preflight_rollout() + + assert status["policy_ready"] is True + control.execute_trajectory.assert_not_called() + control.cancel_trajectory.assert_not_called() + + +def test_preflight_requires_the_configured_coordinator_task( + make_runtime: RuntimeFactory, +) -> None: + policy = FakePolicy(_action_chunk()) + module, control = make_runtime(policy) + _provide_observation(module) + control.list_tasks.return_value = ["another_task"] + + status = module.preflight_rollout() + + assert status["policy_ready"] is False + assert "missing configured rollout task" in (status["last_error"] or "") + assert policy.config_load_count == 0 + + +@pytest.mark.parametrize( + ("positions", "names", "message"), + [ + ([0.0, 0.0, 0.0, float("nan")], JOINTS, "non-finite positions"), + ([0.0, 0.0, 0.0], JOINTS[:-1], "missing configured joints"), + ], +) +def test_preflight_rejects_invalid_live_joints( + make_runtime: RuntimeFactory, + positions: list[float], + names: list[str], + message: str, +) -> None: + policy = FakePolicy(_action_chunk()) + module, control = make_runtime(policy) + timestamp = time.time() + module._on_color_image( + Image(data=np.zeros((4, 5, 3), dtype=np.uint8), format=ImageFormat.RGB, ts=timestamp) + ) + module._on_joint_state(JointState(ts=timestamp, name=names, position=positions)) + + status = module.preflight_rollout() + + assert status["policy_ready"] is False + assert message in (status["last_error"] or "") + control.execute_trajectory.assert_not_called() + + +@pytest.mark.parametrize( + ("image", "image_format"), + [ + (np.zeros((8, 8, 3), dtype=np.uint8), ImageFormat.RGB), + (np.zeros((4, 5, 3), dtype=np.uint8), ImageFormat.BGR), + ], +) +def test_preflight_requires_exact_live_rgb_contract( + make_runtime: RuntimeFactory, + image: NDArray[np.uint8], + image_format: ImageFormat, +) -> None: + policy = FakePolicy(_action_chunk()) + module, control = make_runtime(policy) + timestamp = time.time() + module._on_color_image(Image(data=image, format=image_format, ts=timestamp)) + module._on_joint_state(JointState(ts=timestamp, name=JOINTS, position=[0.0] * len(JOINTS))) + + status = module.preflight_rollout() + + assert status["policy_ready"] is False + assert "no camera image" in (status["last_error"] or "") + control.execute_trajectory.assert_not_called() + + +@pytest.mark.parametrize( + ("lower", "upper", "message"), + [ + ([-1.0] * 3, [1.0] * 3, "range shape"), + ([-1.0, -1.0, -1.0, float("nan")], [1.0] * 4, "non-finite"), + ([2.0] * 4, [1.0] * 4, "min greater than max"), + ], +) +def test_preflight_rejects_invalid_checkpoint_action_bounds( + make_runtime: RuntimeFactory, + lower: list[float], + upper: list[float], + message: str, +) -> None: + policy = FakePolicy(_action_chunk(), action_lower=lower, action_upper=upper) + module, control = make_runtime(policy) + _provide_observation(module) + + status = module.preflight_rollout() -def test_lerobot_feature_validation_rejects_missing_profile_key() -> None: - config = FakeConfig() - config.input_features = {"observation.state": FakeFeature((7,))} + assert status["policy_ready"] is False + assert message in (status["last_error"] or "") + control.execute_trajectory.assert_not_called() - with pytest.raises(ValueError, match="observation.images.wrist"): - _validate_features(config, OPENYAM_QUEST_IO) # type: ignore[arg-type] +def test_preflight_rejects_unavailable_cuda( + make_runtime: RuntimeFactory, + mocker: pytest_mock.MockerFixture, +) -> None: + policy = FakePolicy(_action_chunk()) + module, control = make_runtime(policy, device="cuda") + _provide_observation(module) + mocker.patch.object(torch.cuda, "is_available", return_value=False) -def test_lerobot_action_bounds_are_extracted_for_common_safety_loop() -> None: - lower, upper = _checkpoint_action_bounds(FakePipeline(), 7) + status = module.preflight_rollout() - assert lower.tolist() == [0.0] * 7 - assert upper.tolist() == [1.0] * 7 + assert status["policy_ready"] is False + assert "CUDA is not available" in (status["last_error"] or "") + control.execute_trajectory.assert_not_called() diff --git a/dimos/imitation/policy/lerobot/python/pyproject.toml b/dimos/imitation/policy/lerobot/python/pyproject.toml index ef489674dd..83c9c7a055 100644 --- a/dimos/imitation/policy/lerobot/python/pyproject.toml +++ b/dimos/imitation/policy/lerobot/python/pyproject.toml @@ -46,10 +46,6 @@ strict = true explicit_package_bases = true mypy_path = "../../../../../" -[[tool.mypy.overrides]] -module = ["dimos", "dimos.*"] -follow_imports = "skip" - [[tool.mypy.overrides]] module = [ "lerobot.configs.policies", diff --git a/dimos/imitation/policy/lerobot/test_module.py b/dimos/imitation/policy/lerobot/test_module.py index 8d4cbcb698..1395336f70 100644 --- a/dimos/imitation/policy/lerobot/test_module.py +++ b/dimos/imitation/policy/lerobot/test_module.py @@ -18,16 +18,15 @@ import pytest from dimos.experimental.isolated_python.module import contract_rpc_names -from dimos.imitation.policy.lerobot.module import LeRobotPolicyConfig, OpenYamLeRobotPolicy +from dimos.imitation.policy.lerobot.module import ( + LeRobotPolicyModule, + LeRobotPolicyModuleConfig, +) -def test_generated_contract_has_profile_ports_and_rpc_surface() -> None: - blueprint = OpenYamLeRobotPolicy.blueprint(artifact="unused", task="test") - streams = {stream.name for stream in blueprint.blueprints[0].streams} - - assert OpenYamLeRobotPolicy.implementation == ("dimos_lerobot.runtime:LeRobotPolicyRuntime") - assert streams == {"button_pressed", "wrist_image", "coordinator_joint_state"} - assert contract_rpc_names(OpenYamLeRobotPolicy) == { +def test_contract_imports_without_runtime_dependencies() -> None: + assert LeRobotPolicyModule.implementation == "dimos_lerobot.runtime:LeRobotPolicyRuntime" + assert contract_rpc_names(LeRobotPolicyModule) == { "preflight_rollout", "rollout_status", "start_rollout", @@ -36,27 +35,63 @@ def test_generated_contract_has_profile_ports_and_rpc_surface() -> None: def test_contract_resolves_sibling_runtime_project() -> None: - module = OpenYamLeRobotPolicy(artifact="unused", task="test task") + module = LeRobotPolicyModule( + policy_path="unused", + task="test task", + joint_names=["joint"], + ) try: assert module.runtime_project == Path(__file__).parent / "python" finally: module.stop() -def test_config_rejects_blank_artifact_and_unknown_button() -> None: - with pytest.raises(ValidationError, match="artifact must not be blank"): - LeRobotPolicyConfig(artifact=" ", task="test") - with pytest.raises(ValidationError, match="unknown Quest button"): - LeRobotPolicyConfig(artifact="checkpoint", task="test", rollout_button="NOPE") +@pytest.mark.parametrize( + ("config", "message"), + [ + ( + { + "policy_path": "checkpoint", + "task": "test task", + "joint_names": ["joint1", "joint1"], + }, + "joint_names must not contain duplicates", + ), + ( + { + "policy_path": " ", + "task": "test task", + "joint_names": ["joint1"], + }, + "policy_path must not be blank", + ), + ( + { + "policy_path": "checkpoint", + "task": "test task", + "joint_names": ["joint1"], + "rollout_button": "NOPE", + }, + "unknown Quest button", + ), + ], +) +def test_config_rejects_ambiguous_names(config: dict[str, object], message: str) -> None: + with pytest.raises(ValidationError, match=message): + LeRobotPolicyModuleConfig.model_validate(config) -def test_existing_relative_artifact_is_resolved( +def test_existing_relative_checkpoint_is_resolved_before_isolation( tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: checkpoint = tmp_path / "checkpoint" checkpoint.mkdir() monkeypatch.chdir(tmp_path) - config = LeRobotPolicyConfig(artifact="checkpoint", task="test task") + config = LeRobotPolicyModuleConfig( + policy_path="checkpoint", + task="test task", + joint_names=["joint1"], + ) - assert config.artifact == str(checkpoint) + assert config.policy_path == str(checkpoint) diff --git a/dimos/imitation/policy/module.py b/dimos/imitation/policy/module.py deleted file mode 100644 index 0bf67486c7..0000000000 --- a/dimos/imitation/policy/module.py +++ /dev/null @@ -1,161 +0,0 @@ -# Copyright 2026 Dimensional Inc. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Profile-driven host contract for isolated policy rollout.""" - -from __future__ import annotations - -from pathlib import Path -from typing import ClassVar, Protocol, TypedDict - -from pydantic import Field, field_validator - -from dimos.control.tasks.trajectory_task.trajectory_task import ( - TrajectoryCancellationResult, - TrajectoryExecutionResult, -) -from dimos.core.core import rpc -from dimos.core.stream import In -from dimos.experimental.isolated_python.module import ( - IsolatedPythonModule, - IsolatedPythonModuleConfig, -) -from dimos.imitation.profile import ImageSource, PolicyIOProfile -from dimos.msgs.sensor_msgs.Image import Image -from dimos.msgs.sensor_msgs.JointState import JointState -from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory -from dimos.spec.utils import Spec -from dimos.teleop.quest.quest_types import BUTTON_ALIASES, Buttons - -POLICY_ROLLOUT_TASK_NAME = "policy_rollout" -POLICY_ROLLOUT_INSTANCE_NAME = "PolicyRolloutModule" - - -class PolicyControlSpec(Spec, Protocol): - """Coordinator operations used by policy rollout.""" - - def execute_trajectory( - self, - trajectory: JointTrajectory, - task_name: str, - ) -> TrajectoryExecutionResult: ... - - def cancel_trajectory(self, task_name: str) -> TrajectoryCancellationResult: ... - - def list_tasks(self) -> list[str]: ... - - -class RolloutStatus(TypedDict): - """Operator-facing state of the configured policy rollout.""" - - active: bool - artifact: str - backend: str | None - task: str - device: str | None - policy_ready: bool - observations_ready: bool - chunks_accepted: int - last_error: str | None - - -class PolicyRolloutConfig(IsolatedPythonModuleConfig): - """Backend-neutral rollout and safety configuration.""" - - artifact: str = Field(min_length=1) - task: str = Field(min_length=1) - device: str | None = None - max_observation_age_s: float = Field(default=0.5, gt=0) - max_execution_horizon_s: float = Field(default=0.5, gt=0) - trajectory_task_name: str = POLICY_ROLLOUT_TASK_NAME - rollout_button: str = "A" - - @field_validator("artifact") - @classmethod - def artifact_must_not_be_blank(cls, artifact: str) -> str: - if not artifact.strip(): - raise ValueError("artifact must not be blank") - path = Path(artifact).expanduser() - return str(path.resolve()) if path.exists() else artifact - - @field_validator("trajectory_task_name") - @classmethod - def trajectory_task_name_must_not_be_blank(cls, name: str) -> str: - if not name.strip(): - raise ValueError("trajectory_task_name must not be blank") - return name - - @field_validator("rollout_button") - @classmethod - def rollout_button_must_be_digital(cls, name: str) -> str: - if BUTTON_ALIASES.get(name, name) not in Buttons.BITS: - raise ValueError(f"unknown Quest button {name!r}") - return name - - -class _PolicyModule(IsolatedPythonModule): - """RPC surface shared by all generated policy module declarations.""" - - config: PolicyRolloutConfig - profile: ClassVar[PolicyIOProfile] - button_pressed: In[Buttons] - _control: PolicyControlSpec - - @rpc - def preflight_rollout(self) -> RolloutStatus: - """Load and validate the policy and live inputs without moving the robot.""" - raise NotImplementedError - - @rpc - def start_rollout(self) -> RolloutStatus: - """Start the configured policy until explicitly stopped or it fails.""" - raise NotImplementedError - - @rpc - def stop_rollout(self) -> RolloutStatus: - """Stop rollout publication and clear the policy action queue.""" - raise NotImplementedError - - @rpc - def rollout_status(self) -> RolloutStatus: - """Return the lifecycle and observation state of the configured policy.""" - raise NotImplementedError - - -def declare_policy_module( - name: str, - module_name: str, - profile: PolicyIOProfile, - config_type: type[PolicyRolloutConfig], - implementation: str, -) -> type[_PolicyModule]: - """Declare a stable importable policy module with profile-shaped ports.""" - annotations: dict[str, object] = {"config": config_type} - for source in profile.observations.values(): - annotations[source.stream] = ( - In[Image] if isinstance(source, ImageSource) else In[JointState] - ) - - return type( - name, - (_PolicyModule,), - { - "__annotations__": annotations, - "__doc__": f"Policy rollout module for the {profile.name!r} profile.", - "__module__": module_name, - "__qualname__": name, - "implementation": implementation, - "profile": profile, - }, - ) diff --git a/dimos/imitation/policy/runtime.py b/dimos/imitation/policy/runtime.py deleted file mode 100644 index c3a4cc604e..0000000000 --- a/dimos/imitation/policy/runtime.py +++ /dev/null @@ -1,485 +0,0 @@ -# Copyright 2026 Dimensional Inc. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Backend-neutral live observation alignment and safe rollout execution.""" - -from __future__ import annotations - -from collections import deque -from collections.abc import Callable -from dataclasses import dataclass -import math -from threading import Condition, Event, RLock, Thread, current_thread -import time -from typing import Any, cast - -import numpy as np -from numpy.typing import NDArray -from reactivex.disposable import Disposable - -from dimos.constants import DEFAULT_THREAD_JOIN_TIMEOUT -from dimos.control.tasks.trajectory_task.trajectory_task import TrajectoryExecutionStatus -from dimos.core.core import rpc -from dimos.core.module import Module -from dimos.imitation.policy.backend import PolicyBackend, PolicyBackendInfo -from dimos.imitation.policy.module import ( - PolicyControlSpec, - PolicyRolloutConfig, - RolloutStatus, - _PolicyModule, -) -from dimos.imitation.profile import ImageSource, JointPositionSource, PolicyIOProfile -from dimos.msgs.sensor_msgs.Image import Image, ImageFormat -from dimos.msgs.sensor_msgs.JointState import JointState -from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory -from dimos.msgs.trajectory_msgs.TrajectoryPoint import TrajectoryPoint -from dimos.teleop.quest.quest_types import BUTTON_ALIASES, Buttons -from dimos.utils.logging_config import setup_logger - -logger = setup_logger() - - -@dataclass(frozen=True) -class _TimedValue: - value: NDArray[Any] - ts: float - - -class _PolicyRuntimeMixin: - """Runtime implementation mixed into a generated backend declaration.""" - - profile: PolicyIOProfile - backend_type: type[PolicyBackend] - config: PolicyRolloutConfig - button_pressed: Any - _control: PolicyControlSpec - register_disposable: Callable[[Any], None] - - def __init__(self, **kwargs: Any) -> None: - super().__init__(**kwargs) - self._lock = RLock() - self._observation_changed = Condition(self._lock) - self._buffers: dict[str, deque[_TimedValue]] = { - key: deque(maxlen=64) for key in self.profile.observations - } - self._backend = self.backend_type(self.config) - self._backend_info: PolicyBackendInfo | None = None - self._stop_event = Event() - self._thread: Thread | None = None - self._chunks_accepted = 0 - self._last_error: str | None = None - self._active = False - - @rpc - def start(self) -> None: - Module.start(cast("Module", self)) - streams = {source.stream for source in self.profile.observations.values()} - for stream_name in streams: - stream = getattr(self, stream_name) - self.register_disposable( - Disposable( - stream.subscribe( - lambda message, name=stream_name: self._on_observation(name, message) - ) - ) - ) - self.register_disposable(Disposable(self.button_pressed.subscribe(self._on_button_pressed))) - - @rpc - def stop(self) -> None: - if not self._stop_policy(): - self._cancel_after_stop_timeout() - Module.stop(cast("Module", self)) - - @rpc - def preflight_rollout(self) -> RolloutStatus: - with self._lock: - if self._active: - self._last_error = "cannot preflight while a policy rollout is active" - return self._status_locked() - try: - self._snapshot_observation(time.time()) - except Exception as exc: - self._backend_info = None - self._last_error = str(exc) - return self._status_locked() - - try: - tasks = set(self._control.list_tasks()) - if self.config.trajectory_task_name not in tasks: - raise RuntimeError( - "ControlCoordinator is missing configured rollout task " - f"{self.config.trajectory_task_name!r}" - ) - backend_info = self._backend_info or self._backend.load(self.profile) - self._validate_backend_info(backend_info) - with self._lock: - self._backend_info = backend_info - self._snapshot_observation(time.time()) - self._last_error = None - return self._status_locked() - except Exception as exc: - with self._lock: - self._backend_info = None - self._last_error = str(exc) - return self._status_locked() - - @rpc - def start_rollout(self) -> RolloutStatus: - with self._lock: - if self._thread is not None and self._thread.is_alive(): - self._last_error = "a policy rollout is already active" - return self._status_locked() - if self._backend_info is None: - self._last_error = "policy preflight has not passed" - return self._status_locked() - try: - self._snapshot_observation(time.time()) - except RuntimeError as exc: - self._last_error = str(exc) - return self._status_locked() - - self._stop_event.clear() - self._chunks_accepted = 0 - self._last_error = None - self._active = True - self._thread = Thread( - target=self._run_rollout, - name="policy-rollout", - daemon=True, - ) - self._thread.start() - return self._status_locked() - - @rpc - def stop_rollout(self) -> RolloutStatus: - if not self._stop_policy(): - self._cancel_after_stop_timeout() - return self.rollout_status() - - @rpc - def rollout_status(self) -> RolloutStatus: - with self._lock: - return self._status_locked() - - def _status_locked(self) -> RolloutStatus: - try: - self._snapshot_observation(time.time()) - observations_ready = True - except RuntimeError: - observations_ready = False - return { - "active": self._active, - "artifact": self.config.artifact, - "backend": self._backend_info.name if self._backend_info is not None else None, - "task": self.config.task, - "device": self.config.device, - "policy_ready": self._backend_info is not None, - "observations_ready": observations_ready, - "chunks_accepted": self._chunks_accepted, - "last_error": self._last_error, - } - - def _on_observation(self, stream_name: str, message: object) -> None: - changed_action_state = False - with self._observation_changed: - for key, source in self.profile.observations.items(): - if source.stream != stream_name: - continue - try: - value, ts = _read_source(source, message) - except (TypeError, ValueError) as exc: - logger.warning( - "Ignoring invalid policy observation", - feature=key, - error=str(exc), - ) - continue - self._buffers[key].append(_TimedValue(value=value, ts=ts)) - changed_action_state = changed_action_state or key == self.profile.action_state_key - if changed_action_state: - self._observation_changed.notify_all() - - def _on_button_pressed(self, buttons: Buttons) -> None: - button = BUTTON_ALIASES.get(self.config.rollout_button, self.config.rollout_button) - if not bool(getattr(buttons, button)): - return - with self._lock: - active = self._active - if active: - self.stop_rollout() - else: - self.start_rollout() - - def _snapshot_observation( - self, - now: float, - ) -> tuple[dict[str, NDArray[Any]], NDArray[np.float32], float]: - anchor_key = self.profile.sync.anchor - anchor_buffer = self._buffers[anchor_key] - if not anchor_buffer: - raise RuntimeError(f"no {anchor_key!r} observation has been received") - anchor = anchor_buffer[-1] - selected = {anchor_key: anchor} - tolerance_s = self.profile.sync.tolerance_ms / 1000.0 - - for key, buffer in self._buffers.items(): - if key == anchor_key: - continue - if not buffer: - raise RuntimeError(f"no {key!r} observation has been received") - nearest = min(buffer, key=lambda item: abs(item.ts - anchor.ts)) - skew = abs(nearest.ts - anchor.ts) - if skew > tolerance_s: - raise RuntimeError( - f"{key!r} is {skew * 1000.0:.1f}ms from anchor {anchor_key!r}; " - f"limit is {self.profile.sync.tolerance_ms:.1f}ms" - ) - selected[key] = nearest - - for key, item in selected.items(): - age = now - item.ts - if age > self.config.max_observation_age_s: - raise RuntimeError(f"{key!r} observation is stale by {age:.2f}s") - - state_item = selected[self.profile.action_state_key] - observations = {key: item.value.copy() for key, item in selected.items()} - return observations, np.asarray(state_item.value, dtype=np.float32), state_item.ts - - def _run_rollout(self) -> None: - info: PolicyBackendInfo | None = None - try: - with self._lock: - info = self._backend_info - if info is None: - raise RuntimeError("policy preflight has not passed") - if self._stop_event.is_set(): - return - self._backend.reset() - execution_steps = self._execution_steps(info) - - while not self._stop_event.is_set(): - with self._lock: - observations, state, state_ts = self._snapshot_observation(time.time()) - action_chunk = np.asarray( - self._backend.predict(observations, self.config.task), - dtype=np.float32, - ) - actions = self._validated_actions(action_chunk, info, execution_steps) - if self._stop_event.is_set(): - break - result = self._control.execute_trajectory( - self._trajectory(state, actions), - task_name=self.config.trajectory_task_name, - ) - if result.status is TrajectoryExecutionStatus.START_STATE_MISMATCH: - self._wait_for_newer_joint_state(state_ts) - continue - if result.status is not TrajectoryExecutionStatus.ACCEPTED: - raise RuntimeError( - result.message or f"trajectory rejected: {result.status.name}" - ) - with self._lock: - self._chunks_accepted += 1 - self._stop_event.wait(execution_steps / self.profile.sync.rate_hz) - except Exception as exc: - with self._lock: - self._last_error = str(exc) - logger.exception("Policy execution stopped", error=str(exc)) - finally: - self._stop_event.set() - cancellation_error = self._cancel_trajectory() - reset_error = self._reset_backend() if info is not None else None - with self._lock: - shutdown_errors = [ - error for error in (cancellation_error, reset_error) if error is not None - ] - if shutdown_errors: - shutdown_error = "; ".join(shutdown_errors) - self._last_error = ( - f"{self._last_error}; {shutdown_error}" - if self._last_error is not None - else shutdown_error - ) - self._active = False - - def _execution_steps(self, info: PolicyBackendInfo) -> int: - horizon_steps = math.floor( - self.config.max_execution_horizon_s * self.profile.sync.rate_hz + 1e-9 - ) - if horizon_steps < 1: - raise ValueError("max execution horizon is shorter than one policy step") - return min(info.chunk_length, info.preferred_execution_steps, horizon_steps) - - def _validated_actions( - self, - action_chunk: NDArray[np.float32], - info: PolicyBackendInfo, - execution_steps: int, - ) -> NDArray[np.float32]: - width = len(self.profile.action.demonstration.joints) - if action_chunk.ndim != 2 or action_chunk.shape[1] != width: - raise RuntimeError( - f"policy returned action chunk shape {action_chunk.shape}, expected (steps, {width})" - ) - if action_chunk.shape[0] < execution_steps: - raise RuntimeError( - f"policy returned {action_chunk.shape[0]} steps, expected at least {execution_steps}" - ) - actions = action_chunk[:execution_steps] - if not np.all(np.isfinite(actions)): - raise RuntimeError("policy returned non-finite joint targets") - if info.action_lower is None: - return actions - assert info.action_upper is not None - bounded = np.clip(actions, info.action_lower, info.action_upper) - if np.any(actions != bounded): - logger.warning("Clipped policy actions to backend range") - return bounded - - def _validate_backend_info(self, info: PolicyBackendInfo) -> None: - if info.chunk_length <= 0 or info.preferred_execution_steps <= 0: - raise ValueError("backend chunk and execution step counts must be positive") - bounds = (info.action_lower, info.action_upper) - if (bounds[0] is None) != (bounds[1] is None): - raise ValueError("backend must provide both action bounds or neither") - if bounds[0] is None: - return - assert bounds[1] is not None - shape = (len(self.profile.action.demonstration.joints),) - if bounds[0].shape != shape or bounds[1].shape != shape: - raise ValueError(f"backend action bounds must have shape {shape}") - if not np.all(np.isfinite(bounds[0])) or not np.all(np.isfinite(bounds[1])): - raise ValueError("backend action bounds contain non-finite values") - if np.any(bounds[0] > bounds[1]): - raise ValueError("backend action lower bound exceeds upper bound") - - def _stop_policy(self) -> bool: - with self._lock: - thread = self._thread - self._stop_event.set() - self._observation_changed.notify_all() - if thread is not None and thread is not current_thread(): - thread.join(timeout=DEFAULT_THREAD_JOIN_TIMEOUT) - return thread is None or not thread.is_alive() - - def _cancel_after_stop_timeout(self) -> None: - timeout_error = f"policy rollout did not stop within {DEFAULT_THREAD_JOIN_TIMEOUT} seconds" - cancellation_error = self._cancel_trajectory() - with self._lock: - self._last_error = ( - f"{timeout_error}; {cancellation_error}" - if cancellation_error is not None - else timeout_error - ) - - def _trajectory( - self, - state: NDArray[np.float32], - actions: NDArray[np.float32], - ) -> JointTrajectory: - joints = list(self.profile.action.demonstration.joints) - zeros = [0.0] * len(joints) - points = [ - TrajectoryPoint( - positions=[float(value) for value in state], - velocities=zeros, - time_from_start=0.0, - ) - ] - points.extend( - TrajectoryPoint( - positions=[float(value) for value in action], - velocities=zeros, - time_from_start=(index + 1) / self.profile.sync.rate_hz, - ) - for index, action in enumerate(actions) - ) - return JointTrajectory(joint_names=joints, points=points) - - def _wait_for_newer_joint_state(self, previous_ts: float) -> None: - state_buffer = self._buffers[self.profile.action_state_key] - with self._observation_changed: - self._observation_changed.wait_for( - lambda: self._stop_event.is_set() - or (bool(state_buffer) and state_buffer[-1].ts > previous_ts) - ) - - def _cancel_trajectory(self) -> str | None: - try: - result = self._control.cancel_trajectory(task_name=self.config.trajectory_task_name) - except Exception as exc: - logger.exception( - "Failed to cancel policy trajectory", - task_name=self.config.trajectory_task_name, - ) - return f"Failed to cancel policy trajectory: {exc}" - if result.safe: - return None - return result.message or "Policy trajectory cancellation was uncertain" - - def _reset_backend(self) -> str | None: - try: - self._backend.reset() - except Exception as exc: - logger.exception("Failed to reset policy backend") - return f"Failed to reset policy backend: {exc}" - return None - - -def declare_policy_runtime( - name: str, - module_name: str, - declaration: type[_PolicyModule], - backend_type: type[PolicyBackend], -) -> type[_PolicyModule]: - """Declare an importable runtime subclass using the common rollout loop.""" - return type( - name, - (_PolicyRuntimeMixin, declaration), - { - "__module__": module_name, - "__qualname__": name, - "backend_type": backend_type, - }, - ) - - -def _read_source( - source: ImageSource | JointPositionSource, - message: object, -) -> tuple[NDArray[Any], float]: - if isinstance(source, ImageSource): - if not isinstance(message, Image): - raise TypeError(f"expected Image, got {type(message).__name__}") - if message.format != ImageFormat.RGB or message.data.dtype != np.uint8: - raise ValueError("image must be uint8 RGB") - if message.data.shape != source.shape: - raise ValueError(f"image shape {message.data.shape} does not match {source.shape}") - return np.ascontiguousarray(message.data), message.ts - - if not isinstance(message, JointState): - raise TypeError(f"expected JointState, got {type(message).__name__}") - if len(message.name) != len(message.position): - raise ValueError("JointState names and positions have different lengths") - if len(message.name) != len(set(message.name)): - raise ValueError("JointState contains duplicate joint names") - positions = dict(zip(message.name, message.position, strict=True)) - missing = [joint for joint in source.joints if joint not in positions] - if missing: - raise ValueError(f"JointState is missing configured joints: {missing}") - value = np.asarray([positions[joint] for joint in source.joints], dtype=np.float32) - if not np.all(np.isfinite(value)): - raise ValueError("JointState contains non-finite positions") - return value, message.ts diff --git a/dimos/imitation/policy/test_runtime.py b/dimos/imitation/policy/test_runtime.py deleted file mode 100644 index 1cb964d005..0000000000 --- a/dimos/imitation/policy/test_runtime.py +++ /dev/null @@ -1,320 +0,0 @@ -# Copyright 2026 Dimensional Inc. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -from collections.abc import Iterator, Mapping -from pathlib import Path -import time -from typing import Any - -import numpy as np -from numpy.typing import NDArray -import pytest -import pytest_mock - -from dimos.control.tasks.trajectory_task.trajectory_task import ( - TrajectoryCancellationResult, - TrajectoryCancellationStatus, - TrajectoryExecutionResult, - TrajectoryExecutionStatus, -) -from dimos.imitation.dataprep.core import SyncConfig -from dimos.imitation.policy.backend import PolicyBackendInfo -from dimos.imitation.policy.module import PolicyRolloutConfig, declare_policy_module -from dimos.imitation.policy.runtime import declare_policy_runtime -from dimos.imitation.profile import ( - ImageSource, - JointPositionAction, - JointPositionSource, - PolicyIOProfile, -) -from dimos.msgs.sensor_msgs.Image import Image, ImageFormat -from dimos.msgs.sensor_msgs.JointState import JointState -from dimos.protocol.rpc.pubsubrpc import LCMRPC -from dimos.utils.testing.waiting import wait_until - -JOINTS = ("left", "right") -PROFILE = PolicyIOProfile( - name="runtime-test", - robot_type="test", - observations={ - "top": ImageSource(stream="top_image", shape=(4, 5, 3)), - "left": ImageSource(stream="left_image", shape=(4, 5, 3)), - "state": JointPositionSource(stream="joint_state", joints=JOINTS), - }, - action=JointPositionAction( - key="actions", - demonstration=JointPositionSource(stream="joint_command", joints=JOINTS), - ), - sync=SyncConfig(anchor="top", rate_hz=30.0, tolerance_ms=20.0), -) - - -class _TestConfig(PolicyRolloutConfig): - pass - - -PolicyDeclaration = declare_policy_module( - "RuntimeTestPolicy", - __name__, - PROFILE, - _TestConfig, - "unused:TestRuntime", -) - - -class FakeBackend: - actions = np.arange(60, dtype=np.float32).reshape(30, 2) - - def __init__(self, _config: _TestConfig) -> None: - self.actions = type(self).actions.copy() - self.info = PolicyBackendInfo( - name="fake", - chunk_length=30, - preferred_execution_steps=20, - ) - self.load_count = 0 - self.reset_count = 0 - self.reset_error: str | None = None - - def load(self, _profile: PolicyIOProfile) -> PolicyBackendInfo: - self.load_count += 1 - return self.info - - def reset(self) -> None: - self.reset_count += 1 - if self.reset_error is not None: - raise RuntimeError(self.reset_error) - - def predict( - self, - _observations: Mapping[str, NDArray[Any]], - _task: str, - ) -> NDArray[np.float32]: - return self.actions.copy() - - -RuntimeModule = declare_policy_runtime( - "RuntimeTestModule", - __name__, - PolicyDeclaration, - FakeBackend, -) - - -@pytest.fixture -def runtime( - tmp_path: Path, - mocker: pytest_mock.MockerFixture, -) -> Iterator[tuple[RuntimeModule, Any]]: - mocker.patch("dimos.core.module.get_loop", return_value=(mocker.MagicMock(), None)) - mocker.patch.object(LCMRPC, "__init__", return_value=None) - mocker.patch.object(LCMRPC, "serve_module_rpc", return_value=None) - mocker.patch.object(LCMRPC, "start", return_value=None) - mocker.patch.object(LCMRPC, "stop", return_value=None) - module = RuntimeModule( - _isolated_python_runtime=True, - artifact=str(tmp_path / "artifact"), - task="test task", - ) - control = mocker.MagicMock() - control.list_tasks.return_value = ["policy_rollout"] - control.execute_trajectory.return_value = TrajectoryExecutionResult( - TrajectoryExecutionStatus.ACCEPTED - ) - control.cancel_trajectory.return_value = TrajectoryCancellationResult( - TrajectoryCancellationStatus.ALREADY_STOPPED - ) - mocker.patch.object(module, "_control", control, create=True) - yield module, control - module.stop() - - -def _image(ts: float) -> Image: - return Image(data=np.zeros((4, 5, 3), dtype=np.uint8), format=ImageFormat.RGB, ts=ts) - - -def _provide(module: RuntimeModule, *, top_ts: float, other_ts: float) -> None: - module._on_observation("top_image", _image(top_ts)) - module._on_observation("left_image", _image(other_ts)) - module._on_observation( - "joint_state", - JointState(ts=other_ts, name=list(reversed(JOINTS)), position=[2.0, 1.0]), - ) - - -def test_preflight_requires_every_camera_within_profile_tolerance( - runtime: tuple[RuntimeModule, Any], -) -> None: - module, _control = runtime - now = time.time() - _provide(module, top_ts=now, other_ts=now - 0.021) - - status = module.preflight_rollout() - - assert status["policy_ready"] is False - assert status["observations_ready"] is False - assert "21.0ms" in (status["last_error"] or "") - - -def test_rollout_caps_execution_horizon_and_preserves_profile_joint_order( - runtime: tuple[RuntimeModule, Any], -) -> None: - module, control = runtime - now = time.time() - _provide(module, top_ts=now, other_ts=now - 0.005) - assert module.preflight_rollout()["policy_ready"] is True - - module.start_rollout() - wait_until(lambda: control.execute_trajectory.call_count >= 1, timeout=1.0) - module.stop_rollout() - - trajectory = control.execute_trajectory.call_args.args[0] - assert trajectory.joint_names == list(JOINTS) - assert len(trajectory.points) == 16 - assert trajectory.points[-1].time_from_start == pytest.approx(0.5) - assert trajectory.points[0].positions == [1.0, 2.0] - np.testing.assert_array_equal(trajectory.points[-1].positions, FakeBackend.actions[14]) - - -def test_missing_declared_top_camera_fails_before_backend_load( - runtime: tuple[RuntimeModule, Any], -) -> None: - module, _control = runtime - now = time.time() - module._on_observation("left_image", _image(now)) - module._on_observation( - "joint_state", - JointState(ts=now, name=list(JOINTS), position=[1.0, 2.0]), - ) - - status = module.preflight_rollout() - - assert status["policy_ready"] is False - assert status["last_error"] == "no 'top' observation has been received" - assert module._backend.load_count == 0 - - -def test_start_requires_successful_preflight( - runtime: tuple[RuntimeModule, Any], -) -> None: - module, control = runtime - now = time.time() - _provide(module, top_ts=now, other_ts=now) - - status = module.start_rollout() - - assert status["active"] is False - assert status["last_error"] == "policy preflight has not passed" - control.execute_trajectory.assert_not_called() - - -def test_preflight_requires_configured_coordinator_task_without_loading_backend( - runtime: tuple[RuntimeModule, Any], -) -> None: - module, control = runtime - now = time.time() - _provide(module, top_ts=now, other_ts=now) - control.list_tasks.return_value = ["another_task"] - - status = module.preflight_rollout() - - assert status["policy_ready"] is False - assert "missing configured rollout task" in (status["last_error"] or "") - assert module._backend.load_count == 0 - control.execute_trajectory.assert_not_called() - - -def test_policy_actions_are_clipped_to_backend_bounds( - runtime: tuple[RuntimeModule, Any], -) -> None: - module, control = runtime - module._backend.actions[:, 1] = 2.0 - module._backend.info = PolicyBackendInfo( - name="fake", - chunk_length=30, - preferred_execution_steps=1, - action_lower=np.asarray([-100.0, 0.0], dtype=np.float32), - action_upper=np.asarray([100.0, 1.0], dtype=np.float32), - ) - now = time.time() - _provide(module, top_ts=now, other_ts=now) - - assert module.preflight_rollout()["policy_ready"] is True - module.start_rollout() - wait_until(lambda: control.execute_trajectory.call_count >= 1, timeout=1.0) - module.stop_rollout() - - trajectory = control.execute_trajectory.call_args.args[0] - assert trajectory.points[1].positions[1] == 1.0 - - -@pytest.mark.parametrize( - ("actions", "message"), - [ - (np.zeros((30, 1), dtype=np.float32), "expected (steps, 2)"), - (np.full((30, 2), np.nan, dtype=np.float32), "non-finite joint targets"), - ], -) -def test_invalid_action_chunk_cancels_and_latches_rollout_off( - runtime: tuple[RuntimeModule, Any], - actions: NDArray[np.float32], - message: str, -) -> None: - module, control = runtime - module._backend.actions = actions - now = time.time() - _provide(module, top_ts=now, other_ts=now) - - assert module.preflight_rollout()["policy_ready"] is True - module.start_rollout() - - wait_until(lambda: module.rollout_status()["active"] is False, timeout=1.0) - assert message in (module.rollout_status()["last_error"] or "") - control.execute_trajectory.assert_not_called() - control.cancel_trajectory.assert_called_with(task_name="policy_rollout") - - -def test_trajectory_rejection_cancels_and_latches_rollout_off( - runtime: tuple[RuntimeModule, Any], -) -> None: - module, control = runtime - now = time.time() - _provide(module, top_ts=now, other_ts=now) - control.execute_trajectory.return_value = TrajectoryExecutionResult( - TrajectoryExecutionStatus.POSITION_LIMIT_VIOLATION, - "outside hardware limits", - ) - - assert module.preflight_rollout()["policy_ready"] is True - module.start_rollout() - - wait_until(lambda: module.rollout_status()["active"] is False, timeout=1.0) - assert module.rollout_status()["last_error"] == "outside hardware limits" - control.cancel_trajectory.assert_called_with(task_name="policy_rollout") - - -def test_backend_reset_failure_does_not_leave_rollout_active( - runtime: tuple[RuntimeModule, Any], -) -> None: - module, control = runtime - now = time.time() - _provide(module, top_ts=now, other_ts=now) - assert module.preflight_rollout()["policy_ready"] is True - module._backend.reset_error = "backend state is stuck" - - module.start_rollout() - - wait_until(lambda: module.rollout_status()["active"] is False, timeout=1.0) - assert "backend state is stuck" in (module.rollout_status()["last_error"] or "") - control.cancel_trajectory.assert_called_with(task_name="policy_rollout") diff --git a/dimos/imitation/profile.py b/dimos/imitation/profile.py deleted file mode 100644 index 2b2261d852..0000000000 --- a/dimos/imitation/profile.py +++ /dev/null @@ -1,159 +0,0 @@ -# Copyright 2026 Dimensional Inc. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Typed observation and action contracts shared by collection and rollout.""" - -from __future__ import annotations - -import keyword - -from pydantic import Field, model_validator - -from dimos.imitation.dataprep.core import ( - DataPrepConfig, - FeatureSpec, - OutputConfig, - QualityConfig, - SyncConfig, -) -from dimos.protocol.service.spec import BaseConfig - - -class ImageSource(BaseConfig): - """One RGB image stream consumed under a backend feature key.""" - - stream: str - shape: tuple[int, int, int] - - @model_validator(mode="after") - def validate_source(self) -> ImageSource: - _validate_stream_name(self.stream) - if any(dimension <= 0 for dimension in self.shape): - raise ValueError("image shape must contain positive dimensions") - if self.shape[2] != 3: - raise ValueError("image source must be RGB HWC with three channels") - return self - - -class JointPositionSource(BaseConfig): - """One named JointState position projection.""" - - stream: str - joints: tuple[str, ...] = Field(min_length=1) - - @model_validator(mode="after") - def validate_source(self) -> JointPositionSource: - _validate_stream_name(self.stream) - if any(not joint.strip() for joint in self.joints): - raise ValueError("joint names must not be blank") - if len(set(self.joints)) != len(self.joints): - raise ValueError("joint names must be unique") - return self - - -PolicySource = ImageSource | JointPositionSource - - -class JointPositionAction(BaseConfig): - """Backend action key and the demonstration stream that teaches it.""" - - key: str = Field(min_length=1) - demonstration: JointPositionSource - - -class PolicyIOProfile(BaseConfig): - """Complete feature-key to DimOS-stream contract for one robot setup.""" - - name: str = Field(min_length=1) - robot_type: str = Field(min_length=1) - observations: dict[str, PolicySource] = Field(min_length=1) - action: JointPositionAction - sync: SyncConfig - quality: QualityConfig = QualityConfig() - - @model_validator(mode="after") - def validate_contract(self) -> PolicyIOProfile: - if any(not key.strip() for key in self.observations): - raise ValueError("observation feature keys must not be blank") - if self.action.key in self.observations: - raise ValueError("action key must not also be an observation key") - if self.sync.anchor not in self.observations: - raise ValueError("sync anchor must name an observation feature") - - types_by_stream: dict[str, type[PolicySource]] = {} - sources = [*self.observations.values(), self.action.demonstration] - for source in sources: - existing = types_by_stream.setdefault(source.stream, type(source)) - if existing is not type(source): - raise ValueError( - f"stream {source.stream!r} is declared with conflicting source types" - ) - matching_states = [ - key - for key, source in self.observations.items() - if isinstance(source, JointPositionSource) - and source.joints == self.action.demonstration.joints - ] - if len(matching_states) != 1: - raise ValueError( - "profile must have exactly one joint observation matching the action joints" - ) - return self - - @property - def action_state_key(self) -> str: - """Return the observation feature used as a trajectory's current state.""" - return next( - key - for key, source in self.observations.items() - if isinstance(source, JointPositionSource) - and source.joints == self.action.demonstration.joints - ) - - def dataprep_config(self, *, source: str = "", output: OutputConfig) -> DataPrepConfig: - """Project this live contract into the native-recording dataset schema.""" - observations = { - key: _feature_spec(source_spec) for key, source_spec in self.observations.items() - } - return DataPrepConfig( - source=source, - observation=observations, - action={self.action.key: _feature_spec(self.action.demonstration)}, - sync=self.sync, - quality=self.quality, - output=output, - ) - - -def _feature_spec(source: PolicySource) -> FeatureSpec: - if isinstance(source, ImageSource): - return FeatureSpec( - stream=source.stream, - field="data", - dtype="video", - shape=source.shape, - names=["height", "width", "channels"], - ) - return FeatureSpec( - stream=source.stream, - field="position", - dtype="float32", - shape=(len(source.joints),), - names=list(source.joints), - ) - - -def _validate_stream_name(name: str) -> None: - if not name.isidentifier() or keyword.iskeyword(name): - raise ValueError(f"stream {name!r} must be a Python identifier") diff --git a/dimos/imitation/test_profile.py b/dimos/imitation/test_profile.py deleted file mode 100644 index d5d6a7b330..0000000000 --- a/dimos/imitation/test_profile.py +++ /dev/null @@ -1,107 +0,0 @@ -# Copyright 2026 Dimensional Inc. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -from pathlib import Path - -from pydantic import ValidationError -import pytest - -from dimos.imitation.dataprep.core import OutputConfig, SyncConfig -from dimos.imitation.profile import ( - ImageSource, - JointPositionAction, - JointPositionSource, - PolicyIOProfile, -) - - -def _profile() -> PolicyIOProfile: - joints = ("left_joint", "right_joint") - return PolicyIOProfile( - name="dual-test", - robot_type="dual_test", - observations={ - "left": ImageSource(stream="left_image", shape=(480, 640, 3)), - "state": JointPositionSource(stream="joint_state", joints=joints), - }, - action=JointPositionAction( - key="actions", - demonstration=JointPositionSource(stream="joint_command", joints=joints), - ), - sync=SyncConfig(anchor="left", rate_hz=30.0, tolerance_ms=20.0), - ) - - -def test_profile_builds_matching_dataprep_config(tmp_path: Path) -> None: - profile = _profile() - output = OutputConfig(path=tmp_path / "dataset") - - config = profile.dataprep_config(source="recording.mcap", output=output) - - assert config.source == "recording.mcap" - assert config.output == output - assert config.observation["left"].stream == "left_image" - assert config.observation["left"].shape == (480, 640, 3) - assert config.observation["state"].names == ["left_joint", "right_joint"] - assert config.action["actions"].stream == "joint_command" - assert config.sync.anchor == "left" - - -@pytest.mark.parametrize( - ("update", "message"), - [ - ({"sync": SyncConfig(anchor="missing", rate_hz=30.0, tolerance_ms=20.0)}, "anchor"), - ( - { - "action": JointPositionAction( - key="left", - demonstration=JointPositionSource(stream="joint_command", joints=("joint",)), - ) - }, - "action key", - ), - ], -) -def test_profile_rejects_ambiguous_feature_contracts( - update: dict[str, object], message: str -) -> None: - values = _profile().model_dump() - values.update(update) - - with pytest.raises(ValidationError, match=message): - PolicyIOProfile.model_validate(values) - - -def test_profile_rejects_conflicting_types_on_one_stream() -> None: - values = _profile().model_dump() - values["action"] = { - "key": "actions", - "demonstration": {"stream": "left_image", "joints": ["joint"]}, - } - - with pytest.raises(ValidationError, match="conflicting source types"): - PolicyIOProfile.model_validate(values) - - -@pytest.mark.parametrize( - "source", - [ - {"stream": "left-image", "shape": (480, 640, 3)}, - {"stream": "left_image", "shape": (480, 640, 1)}, - {"stream": "left_image", "shape": (0, 640, 3)}, - ], -) -def test_image_source_requires_an_rgb_hwc_python_port(source: dict[str, object]) -> None: - with pytest.raises(ValidationError): - ImageSource.model_validate(source) diff --git a/dimos/robot/all_blueprints.py b/dimos/robot/all_blueprints.py index bb9cd4b0d3..f62592d68e 100644 --- a/dimos/robot/all_blueprints.py +++ b/dimos/robot/all_blueprints.py @@ -255,6 +255,7 @@ "joystick-module": "dimos.robot.unitree.b1.joystick_module.JoystickModule", "keyboard-teleop": "dimos.robot.unitree.keyboard_teleop.KeyboardTeleop", "keyboard-teleop-module": "dimos.teleop.keyboard.keyboard_teleop_module.KeyboardTeleopModule", + "le-robot-policy-module": "dimos.imitation.policy.lerobot.module.LeRobotPolicyModule", "lidar-window-relocalization": "dimos.mapping.relocalization.lidar.module.LidarWindowRelocalization", "local-map-relocalization": "dimos.mapping.relocalization.lidar.module.LocalMapRelocalization", "m20-camera-relay": "dimos.robot.deeprobotics.m20.camera.M20CameraRelay", diff --git a/dimos/robot/manipulators/dual_openyam/blueprints/learning_collection.py b/dimos/robot/manipulators/dual_openyam/blueprints/learning_collection.py index 57b7806750..4307f70430 100644 --- a/dimos/robot/manipulators/dual_openyam/blueprints/learning_collection.py +++ b/dimos/robot/manipulators/dual_openyam/blueprints/learning_collection.py @@ -12,21 +12,16 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Dual OpenYAM Quest collection with two independently declared cameras.""" +"""Dual OpenYAM Quest collection with profile-defined camera inputs.""" from pathlib import Path from dimos.core.coordination.blueprints import Blueprint, autoconnect -from dimos.experimental.memory.rust_recorder import RustMcapStoreConfig from dimos.imitation.cameras import CameraDevice, profile_cameras from dimos.imitation.collection.episode_monitor import EpisodeMonitorModule -from dimos.robot.manipulators.dual_openyam.blueprints.teleop import ( - build_dual_openyam_quest_teleop, -) -from dimos.robot.manipulators.dual_openyam.learning import ( - DUAL_OPENYAM_TWO_WRIST_IO, - DualOpenYamQuestRecorder, -) +from dimos.imitation.collection.native_recorder import collection_recorder +from dimos.imitation.collection.profile import CollectionProfile +from dimos.robot.manipulators.dual_openyam.learning import DUAL_OPENYAM_COLLECTION def build_dual_openyam_quest_collection( @@ -34,23 +29,23 @@ def build_dual_openyam_quest_collection( recording: Path, task: str, cameras: dict[str, CameraDevice], + profile: CollectionProfile = DUAL_OPENYAM_COLLECTION, left_can_port: str | None = None, right_can_port: str | None = None, ) -> Blueprint: - """Build a bimanual Quest collection session and two wrist cameras.""" - camera_blueprints, camera_remappings = profile_cameras( - DUAL_OPENYAM_TWO_WRIST_IO, - cameras, + """Build bimanual collection; camera count and feature names come from the profile.""" + camera_blueprints, remappings = profile_cameras(profile, cameras) + # Pink and the dual robot model are optional until this stack is selected. + from dimos.robot.manipulators.dual_openyam.blueprints.teleop import ( + build_dual_openyam_quest_teleop, ) + return autoconnect( - DualOpenYamQuestRecorder.blueprint( - store=RustMcapStoreConfig(path=str(recording)), - record_tf=False, - ), + collection_recorder(profile=profile, recording=recording), EpisodeMonitorModule.blueprint(task=task), build_dual_openyam_quest_teleop( left_can_port=left_can_port, right_can_port=right_can_port, ), *camera_blueprints, - ).remappings(camera_remappings) + ).remappings(remappings) diff --git a/dimos/robot/manipulators/dual_openyam/blueprints/learning_rollout.py b/dimos/robot/manipulators/dual_openyam/blueprints/learning_rollout.py deleted file mode 100644 index 0d2a40fc3d..0000000000 --- a/dimos/robot/manipulators/dual_openyam/blueprints/learning_rollout.py +++ /dev/null @@ -1,78 +0,0 @@ -# Copyright 2026 Dimensional Inc. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Released ABC-DiT rollout for the complete Dual OpenYAM entity.""" - -from dimos.constants import DEFAULT_CAPACITY_COLOR_IMAGE -from dimos.control.coordinator import TaskConfig -from dimos.core.coordination.blueprints import Blueprint, autoconnect -from dimos.core.transport import pSHMTransport -from dimos.imitation.cameras import CameraDevice, profile_cameras -from dimos.imitation.policy.abc.module import DualOpenYamAbcPolicy -from dimos.imitation.policy.module import ( - POLICY_ROLLOUT_INSTANCE_NAME, - POLICY_ROLLOUT_TASK_NAME, -) -from dimos.msgs.sensor_msgs.Image import Image -from dimos.robot.manipulators.dual_openyam.blueprints.basic import DualOpenYamCoordinator -from dimos.robot.manipulators.dual_openyam.learning import ABC_JOINTS, DUAL_OPENYAM_ABC_IO - - -def build_dual_openyam_abc_rollout( - *, - artifact: str, - task: str, - cameras: dict[str, CameraDevice], - device: str | None = None, - quest_control: bool = False, - left_can_port: str | None = None, - right_can_port: str | None = None, -) -> Blueprint: - """Build three-camera ABC rollout; a synthetic top view is never accepted.""" - if quest_control: - raise ValueError("Quest takeover is not defined for Dual OpenYAM ABC rollout") - camera_blueprints, camera_remappings = profile_cameras(DUAL_OPENYAM_ABC_IO, cameras) - blueprint = autoconnect( - DualOpenYamAbcPolicy.blueprint( - instance_name=POLICY_ROLLOUT_INSTANCE_NAME, - artifact=artifact, - task=task, - device=device, - trajectory_task_name=POLICY_ROLLOUT_TASK_NAME, - ), - DualOpenYamCoordinator.blueprint( - instance_name="ControlCoordinator", - left_can_port=left_can_port, - right_can_port=right_can_port, - tasks=[ - TaskConfig( - name=POLICY_ROLLOUT_TASK_NAME, - type="trajectory", - joint_names=list(ABC_JOINTS), - priority=10, - params={"start_position_tolerance": 0.05}, - ) - ], - ), - *camera_blueprints, - ).remappings(camera_remappings) - return blueprint.transports( - { - (stream, Image): pSHMTransport.spec( - f"/{stream}", - default_capacity=DEFAULT_CAPACITY_COLOR_IMAGE, - ) - for stream in ("top_image", "left_wrist_image", "right_wrist_image") - } - ) diff --git a/dimos/robot/manipulators/dual_openyam/blueprints/test_learning.py b/dimos/robot/manipulators/dual_openyam/blueprints/test_learning.py index 6fd93478e9..595e8dd161 100644 --- a/dimos/robot/manipulators/dual_openyam/blueprints/test_learning.py +++ b/dimos/robot/manipulators/dual_openyam/blueprints/test_learning.py @@ -12,23 +12,18 @@ # See the License for the specific language governing permissions and # limitations under the License. -from pathlib import Path import pytest +from dimos.experimental.memory.rust_recorder import RustRecorder from dimos.hardware.sensors.camera.module import CameraModule -from dimos.imitation.policy.abc.module import DualOpenYamAbcPolicy -from dimos.imitation.policy.module import POLICY_ROLLOUT_INSTANCE_NAME, POLICY_ROLLOUT_TASK_NAME from dimos.robot.manipulators.dual_openyam.blueprints.learning_collection import ( build_dual_openyam_quest_collection, ) -from dimos.robot.manipulators.dual_openyam.blueprints.learning_rollout import ( - build_dual_openyam_abc_rollout, -) -from dimos.robot.manipulators.dual_openyam.learning import ABC_JOINTS, DualOpenYamQuestRecorder +from dimos.robot.manipulators.dual_openyam.learning import DUAL_OPENYAM_COLLECTION -def test_dual_collection_declares_two_distinct_camera_sources(tmp_path: Path) -> None: +def test_dual_collection_declares_two_distinct_cameras_and_both_buses(tmp_path): blueprint = build_dual_openyam_quest_collection( recording=tmp_path / "dual.mcap", task="fold towel", @@ -37,14 +32,15 @@ def test_dual_collection_declares_two_distinct_camera_sources(tmp_path: Path) -> right_can_port="follower_r", ) cameras = [atom for atom in blueprint.active_blueprints if atom.module is CameraModule] - - assert any(atom.module is DualOpenYamQuestRecorder for atom in blueprint.active_blueprints) + assert issubclass(blueprint.active_blueprints[0].module, RustRecorder) assert [camera.kwargs["hardware"].camera_index for camera in cameras] == [0, 1] - assert blueprint.remapping_map[("PolicyCamera_left_wrist_image", "color_image")] == ( - "left_wrist_image" + assert ( + blueprint.remapping_map[("CollectionCamera_left_wrist_image", "color_image")] + == "left_wrist_image" ) - assert blueprint.remapping_map[("PolicyCamera_right_wrist_image", "color_image")] == ( - "right_wrist_image" + assert ( + blueprint.remapping_map[("CollectionCamera_right_wrist_image", "color_image")] + == "right_wrist_image" ) coordinator = next( atom for atom in blueprint.active_blueprints if atom.name == "ControlCoordinator" @@ -53,34 +49,40 @@ def test_dual_collection_declares_two_distinct_camera_sources(tmp_path: Path) -> assert coordinator.kwargs["right_can_port"] == "follower_r" -def test_abc_rollout_requires_all_three_physical_cameras() -> None: - with pytest.raises(ValueError, match="top_image"): - build_dual_openyam_abc_rollout( - artifact="checkpoint.pt", - task="put bottles in bin", - cameras={"left_wrist_image": 0, "right_wrist_image": 1}, - ) - - -def test_abc_rollout_uses_released_action_order_and_stable_rpc_name() -> None: - blueprint = build_dual_openyam_abc_rollout( - artifact="checkpoint.pt", - task="put bottles in bin", - cameras={"top_image": 0, "left_wrist_image": 1, "right_wrist_image": 2}, - left_can_port="follower_l", - right_can_port="follower_r", +def test_custom_profile_adds_overhead_camera_without_a_new_recorder(tmp_path): + profile = DUAL_OPENYAM_COLLECTION.model_copy(deep=True) + profile.observations["overhead"] = profile.observations[ + "observation.images.left_wrist" + ].model_copy( + update={"stream": "overhead_image"}, ) - policy = next( - atom for atom in blueprint.active_blueprints if atom.module is DualOpenYamAbcPolicy - ) - coordinator = next( - atom for atom in blueprint.active_blueprints if atom.name == "ControlCoordinator" + blueprint = build_dual_openyam_quest_collection( + profile=profile, + recording=tmp_path / "three.mcap", + task="fold towel", + cameras={"left_wrist_image": 0, "right_wrist_image": 1, "overhead_image": 2}, ) + recorder = blueprint.active_blueprints[0] + assert {s.name for s in recorder.streams} >= { + "left_wrist_image", + "right_wrist_image", + "overhead_image", + "status", + } + assert len([atom for atom in blueprint.active_blueprints if atom.module is CameraModule]) == 3 - assert policy.name == POLICY_ROLLOUT_INSTANCE_NAME - rollout_task = next( - task for task in coordinator.kwargs["tasks"] if task.name == POLICY_ROLLOUT_TASK_NAME - ) - assert rollout_task.joint_names == list(ABC_JOINTS) - assert coordinator.kwargs["left_can_port"] == "follower_l" - assert coordinator.kwargs["right_can_port"] == "follower_r" + +@pytest.mark.parametrize( + ("devices", "error"), + [ + ({"left_wrist_image": 0}, "missing cameras.*right_wrist_image"), + ({"left_wrist_image": 0, "right_wrist_image": 1, "typo": 2}, "unknown cameras.*typo"), + ], +) +def test_invalid_camera_bindings_fail_before_hardware_import(tmp_path, devices, error): + with pytest.raises(ValueError, match=error): + build_dual_openyam_quest_collection( + recording=tmp_path / "dual.mcap", + task="fold", + cameras=devices, + ) diff --git a/dimos/robot/manipulators/dual_openyam/config.py b/dimos/robot/manipulators/dual_openyam/config.py index 0fb28db2b3..d16a323c3b 100644 --- a/dimos/robot/manipulators/dual_openyam/config.py +++ b/dimos/robot/manipulators/dual_openyam/config.py @@ -20,27 +20,22 @@ from dimos.hardware.whole_body.spec import WholeBodyConfig from dimos.manipulation.planning.groups.models import PlanningGroupDefinition from dimos.manipulation.planning.spec.config import RobotModelConfig +from dimos.robot.manipulators.dual_openyam.joints import ( + DUAL_OPENYAM_ARM_JOINTS as DUAL_OPENYAM_ARM_JOINTS, + DUAL_OPENYAM_DOF_PER_ARM as DUAL_OPENYAM_DOF_PER_ARM, + DUAL_OPENYAM_GRIPPER_JOINTS as DUAL_OPENYAM_GRIPPER_JOINTS, + DUAL_OPENYAM_JOINTS as DUAL_OPENYAM_JOINTS, + DUAL_OPENYAM_LEFT_ARM_JOINTS as DUAL_OPENYAM_LEFT_ARM_JOINTS, + DUAL_OPENYAM_RIGHT_ARM_JOINTS as DUAL_OPENYAM_RIGHT_ARM_JOINTS, +) from dimos.robot.manipulators.dual_openyam.model import ( DUAL_OPENYAM_MODEL, ) from dimos.robot.manipulators.openyam.config import OPENYAM_HOME_JOINTS -DUAL_OPENYAM_DOF_PER_ARM = 6 DUAL_OPENYAM_HARDWARE_ID = "dual_openyam" DUAL_OPENYAM_ADAPTER_TYPE = "dual_openyam_damiao" DUAL_OPENYAM_SIDES = ("left", "right") -DUAL_OPENYAM_LEFT_ARM_JOINTS = [ - f"left_joint{index}" for index in range(1, DUAL_OPENYAM_DOF_PER_ARM + 1) -] -DUAL_OPENYAM_RIGHT_ARM_JOINTS = [ - f"right_joint{index}" for index in range(1, DUAL_OPENYAM_DOF_PER_ARM + 1) -] -DUAL_OPENYAM_ARM_JOINTS = [ - *DUAL_OPENYAM_LEFT_ARM_JOINTS, - *DUAL_OPENYAM_RIGHT_ARM_JOINTS, -] -DUAL_OPENYAM_GRIPPER_JOINTS = ["left_arm/gripper", "right_arm/gripper"] -DUAL_OPENYAM_JOINTS = [*DUAL_OPENYAM_ARM_JOINTS, *DUAL_OPENYAM_GRIPPER_JOINTS] DUAL_OPENYAM_HOME_PER_ARM = list(OPENYAM_HOME_JOINTS) DUAL_OPENYAM_HOME_JOINTS = [*DUAL_OPENYAM_HOME_PER_ARM, *DUAL_OPENYAM_HOME_PER_ARM] _ARM_KP = (80.0, 80.0, 80.0, 10.0, 10.0, 10.0) diff --git a/dimos/robot/manipulators/dual_openyam/joints.py b/dimos/robot/manipulators/dual_openyam/joints.py new file mode 100644 index 0000000000..45e1e89b2e --- /dev/null +++ b/dimos/robot/manipulators/dual_openyam/joints.py @@ -0,0 +1,22 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Dual OpenYAM joint names without loading robot models.""" + +DUAL_OPENYAM_DOF_PER_ARM = 6 +DUAL_OPENYAM_LEFT_ARM_JOINTS = [f"left_joint{i}" for i in range(1, DUAL_OPENYAM_DOF_PER_ARM + 1)] +DUAL_OPENYAM_RIGHT_ARM_JOINTS = [f"right_joint{i}" for i in range(1, DUAL_OPENYAM_DOF_PER_ARM + 1)] +DUAL_OPENYAM_ARM_JOINTS = [*DUAL_OPENYAM_LEFT_ARM_JOINTS, *DUAL_OPENYAM_RIGHT_ARM_JOINTS] +DUAL_OPENYAM_GRIPPER_JOINTS = ["left_arm/gripper", "right_arm/gripper"] +DUAL_OPENYAM_JOINTS = [*DUAL_OPENYAM_ARM_JOINTS, *DUAL_OPENYAM_GRIPPER_JOINTS] diff --git a/dimos/robot/manipulators/dual_openyam/learning.py b/dimos/robot/manipulators/dual_openyam/learning.py index 01b8852575..af37ef6e8e 100644 --- a/dimos/robot/manipulators/dual_openyam/learning.py +++ b/dimos/robot/manipulators/dual_openyam/learning.py @@ -12,96 +12,48 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Distinct Dual OpenYAM collection and released-ABC rollout profiles.""" +"""Dual-arm collection preset, independent of policy backends and hardware.""" -from dimos.imitation.collection.native_recorder import declare_recorder +from dimos.imitation.collection.profile import CollectionFeature, CollectionProfile from dimos.imitation.dataprep.core import QualityConfig, SyncConfig -from dimos.imitation.profile import ( - ImageSource, - JointPositionAction, - JointPositionSource, - PolicyIOProfile, -) -from dimos.robot.manipulators.dual_openyam.config import ( - DUAL_OPENYAM_GRIPPER_JOINTS, - DUAL_OPENYAM_JOINTS, - DUAL_OPENYAM_LEFT_ARM_JOINTS, - DUAL_OPENYAM_RIGHT_ARM_JOINTS, -) +from dimos.msgs.sensor_msgs.Image import Image +from dimos.msgs.sensor_msgs.JointState import JointState +from dimos.robot.manipulators.dual_openyam.joints import DUAL_OPENYAM_JOINTS -DUAL_OPENYAM_CAMERA_SHAPE = (480, 640, 3) -DUAL_OPENYAM_FPS = 30.0 -ABC_JOINTS = ( - *DUAL_OPENYAM_LEFT_ARM_JOINTS, - DUAL_OPENYAM_GRIPPER_JOINTS[0], - *DUAL_OPENYAM_RIGHT_ARM_JOINTS, - DUAL_OPENYAM_GRIPPER_JOINTS[1], -) - -_quality = QualityConfig( - mode="strict", - min_source_rate_ratio=0.95, - max_camera_gap_ms=100.0, - max_alignment_error_ms=20.0, -) - -DUAL_OPENYAM_TWO_WRIST_IO = PolicyIOProfile( +DUAL_OPENYAM_COLLECTION = CollectionProfile( name="dual-openyam-quest", robot_type="dual_openyam", observations={ - "observation.images.left_wrist": ImageSource( - stream="left_wrist_image", - shape=DUAL_OPENYAM_CAMERA_SHAPE, - ), - "observation.images.right_wrist": ImageSource( - stream="right_wrist_image", - shape=DUAL_OPENYAM_CAMERA_SHAPE, - ), - "observation.state": JointPositionSource( + **{ + f"observation.images.{side}_wrist": CollectionFeature( + stream=f"{side}_wrist_image", + message_type=Image, + field="data", + dtype="video", + shape=(480, 640, 3), + names=["height", "width", "channels"], + ) + for side in ("left", "right") + }, + "observation.state": CollectionFeature( stream="coordinator_joint_state", - joints=tuple(DUAL_OPENYAM_JOINTS), + message_type=JointState, + field="position", + dtype="float32", + shape=(len(DUAL_OPENYAM_JOINTS),), + names=list(DUAL_OPENYAM_JOINTS), ), }, - action=JointPositionAction( - key="action", - demonstration=JointPositionSource( + actions={ + "action": CollectionFeature( stream="applied_joint_position_command", - joints=tuple(DUAL_OPENYAM_JOINTS), - ), - ), - sync=SyncConfig( - anchor="observation.images.left_wrist", - rate_hz=DUAL_OPENYAM_FPS, - tolerance_ms=20.0, - ), - quality=_quality, -) - -DUAL_OPENYAM_ABC_IO = PolicyIOProfile( - name="dual-openyam-abc", - robot_type="dual_openyam", - observations={ - "top": ImageSource(stream="top_image", shape=DUAL_OPENYAM_CAMERA_SHAPE), - "left": ImageSource(stream="left_wrist_image", shape=DUAL_OPENYAM_CAMERA_SHAPE), - "right": ImageSource(stream="right_wrist_image", shape=DUAL_OPENYAM_CAMERA_SHAPE), - "state": JointPositionSource( - stream="coordinator_joint_state", - joints=ABC_JOINTS, + message_type=JointState, + field="position", + dtype="float32", + shape=(len(DUAL_OPENYAM_JOINTS),), + names=list(DUAL_OPENYAM_JOINTS), ), }, - action=JointPositionAction( - key="actions", - demonstration=JointPositionSource( - stream="applied_joint_position_command", - joints=ABC_JOINTS, - ), - ), - sync=SyncConfig(anchor="top", rate_hz=DUAL_OPENYAM_FPS, tolerance_ms=20.0), - quality=_quality, -) - -DualOpenYamQuestRecorder = declare_recorder( - "DualOpenYamQuestRecorder", - __name__, - DUAL_OPENYAM_TWO_WRIST_IO, + sync=SyncConfig(anchor="observation.images.left_wrist", rate_hz=30, tolerance_ms=20), + quality=QualityConfig(mode="strict"), ) diff --git a/pyproject.toml b/pyproject.toml index 3e36ba43bc..5c6af291ce 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -54,9 +54,6 @@ exclude = [ # shipping arbitrary manifests from the source tree. "imitation/policy/lerobot/python/pyproject.toml", "imitation/policy/lerobot/python/uv.lock", - "imitation/policy/abc/python/pyproject.toml", - "imitation/policy/abc/python/uv.lock", - "imitation/policy/abc/python/VENDORED.md", ] [tool.setuptools.exclude-package-data] @@ -598,7 +595,6 @@ exclude = [ "venv", "libs", "external", - "dimos/imitation/policy/abc/python/abc_minimal", "src" ] @@ -642,7 +638,7 @@ strict = true warn_unused_ignores = false untyped_calls_exclude = ["zenoh"] explicit_package_bases = true -exclude = "^dimos/models/Detic(/|$)|^dimos/imitation/policy/(lerobot|abc)/python/|.*/test_.|.*/tool_.|.*/conftest.py*" +exclude = "^dimos/models/Detic(/|$)|^dimos/imitation/policy/lerobot/python/|.*/test_.|.*/tool_.|.*/conftest.py*" [[tool.mypy.overrides]] module = [ From 46985aa742dcb089567d9f777114b982962d861b Mon Sep 17 00:00:00 2001 From: cc Date: Thu, 10 Sep 2026 13:08:09 -0700 Subject: [PATCH 3/4] refactor(imitation): use blueprint launches and attached controls --- dimos/imitation/cameras.py | 72 -------------- dimos/robot/all_blueprints.py | 1 + .../blueprints/learning_collection.py | 63 ++++++------- .../dual_openyam/blueprints/test_learning.py | 93 ++++++------------- 4 files changed, 60 insertions(+), 169 deletions(-) delete mode 100644 dimos/imitation/cameras.py diff --git a/dimos/imitation/cameras.py b/dimos/imitation/cameras.py deleted file mode 100644 index df63be25aa..0000000000 --- a/dimos/imitation/cameras.py +++ /dev/null @@ -1,72 +0,0 @@ -# Copyright 2026 Dimensional Inc. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Camera blueprints generated from collection image features.""" - -from __future__ import annotations - -from collections.abc import Mapping - -from dimos.core.coordination.blueprints import Blueprint -from dimos.hardware.sensors.camera.module import CameraModule -from dimos.hardware.sensors.camera.webcam import WebcamConfig -from dimos.imitation.collection.profile import CollectionProfile - -CameraDevice = int | str - - -def profile_cameras( - profile: CollectionProfile, - devices: Mapping[str, CameraDevice], -) -> tuple[list[Blueprint], list[tuple[str, str, str]]]: - """Build cameras and explicit output remappings for a collection profile.""" - image_sources = profile.camera_features() - missing = sorted(set(image_sources) - set(devices)) - unknown = sorted(set(devices) - set(image_sources)) - if missing or unknown: - details = [] - if missing: - details.append(f"missing cameras: {missing}") - if unknown: - details.append(f"unknown cameras: {unknown}") - raise ValueError("; ".join(details)) - - blueprints: list[Blueprint] = [] - remappings: list[tuple[str, str, str]] = [] - for stream_name, source in image_sources.items(): - if len(source.shape) != 3 or source.shape[2] != 3: - raise ValueError(f"Camera {stream_name!r} requires an HWC RGB shape") - height, width, _channels = source.shape - instance_name = f"CollectionCamera_{stream_name}" - blueprints.append( - CameraModule.blueprint( - instance_name=instance_name, - hardware=WebcamConfig( - camera_index=devices[stream_name], - width=width, - height=height, - fps=profile.sync.rate_hz, - frame_id_prefix=stream_name, - ), - frame_id=f"{stream_name}_camera_link", - ) - ) - remappings.extend( - [ - (instance_name, "color_image", stream_name), - (instance_name, "camera_info", f"{stream_name}_camera_info"), - (instance_name, "tf", f"{stream_name}_tf"), - ] - ) - return blueprints, remappings diff --git a/dimos/robot/all_blueprints.py b/dimos/robot/all_blueprints.py index f62592d68e..e98e88de6e 100644 --- a/dimos/robot/all_blueprints.py +++ b/dimos/robot/all_blueprints.py @@ -65,6 +65,7 @@ "drone-agentic": "dimos.robot.drone.blueprints.agentic.drone_agentic:drone_agentic", "drone-basic": "dimos.robot.drone.blueprints.basic.drone_basic:drone_basic", "dual-openyam-planner-coordinator": "dimos.robot.manipulators.dual_openyam.blueprints.basic:dual_openyam_planner_coordinator", + "dual-openyam-quest-collection": "dimos.robot.manipulators.dual_openyam.blueprints.learning_collection:dual_openyam_quest_collection", "dual-xarm6-planner-coordinator": "dimos.robot.manipulators.xarm.blueprints.basic:dual_xarm6_planner_coordinator", "go2-zenoh-basic": "dimos.robot.unitree.go2.zenoh.blueprints:go2_zenoh_basic", "go2-zenoh-htc": "dimos.robot.unitree.go2.zenoh.blueprints:go2_zenoh_htc", diff --git a/dimos/robot/manipulators/dual_openyam/blueprints/learning_collection.py b/dimos/robot/manipulators/dual_openyam/blueprints/learning_collection.py index 4307f70430..8cc37225a8 100644 --- a/dimos/robot/manipulators/dual_openyam/blueprints/learning_collection.py +++ b/dimos/robot/manipulators/dual_openyam/blueprints/learning_collection.py @@ -12,40 +12,41 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Dual OpenYAM Quest collection with profile-defined camera inputs.""" +"""Dual OpenYAM collection using ordinary blueprint configuration.""" -from pathlib import Path - -from dimos.core.coordination.blueprints import Blueprint, autoconnect -from dimos.imitation.cameras import CameraDevice, profile_cameras +from dimos.core.coordination.blueprints import autoconnect +from dimos.hardware.sensors.camera.module import CameraModule +from dimos.hardware.sensors.camera.webcam import WebcamConfig from dimos.imitation.collection.episode_monitor import EpisodeMonitorModule from dimos.imitation.collection.native_recorder import collection_recorder -from dimos.imitation.collection.profile import CollectionProfile +from dimos.robot.manipulators.dual_openyam.blueprints.teleop import teleop_quest_dual_openyam from dimos.robot.manipulators.dual_openyam.learning import DUAL_OPENYAM_COLLECTION - -def build_dual_openyam_quest_collection( - *, - recording: Path, - task: str, - cameras: dict[str, CameraDevice], - profile: CollectionProfile = DUAL_OPENYAM_COLLECTION, - left_can_port: str | None = None, - right_can_port: str | None = None, -) -> Blueprint: - """Build bimanual collection; camera count and feature names come from the profile.""" - camera_blueprints, remappings = profile_cameras(profile, cameras) - # Pink and the dual robot model are optional until this stack is selected. - from dimos.robot.manipulators.dual_openyam.blueprints.teleop import ( - build_dual_openyam_quest_teleop, - ) - - return autoconnect( - collection_recorder(profile=profile, recording=recording), - EpisodeMonitorModule.blueprint(task=task), - build_dual_openyam_quest_teleop( - left_can_port=left_can_port, - right_can_port=right_can_port, +dual_openyam_quest_collection = autoconnect( + teleop_quest_dual_openyam, + CameraModule.blueprint( + instance_name="left_wrist", + hardware=WebcamConfig( + camera_index=0, width=640, height=480, fps=30, frame_id_prefix="left_wrist_image" + ), + frame_id="left_wrist_camera_link", + ), + CameraModule.blueprint( + instance_name="right_wrist", + hardware=WebcamConfig( + camera_index=1, width=640, height=480, fps=30, frame_id_prefix="right_wrist_image" ), - *camera_blueprints, - ).remappings(remappings) + frame_id="right_wrist_camera_link", + ), + collection_recorder(profile=DUAL_OPENYAM_COLLECTION), + EpisodeMonitorModule.blueprint(instance_name="episodes"), +).remappings( + [ + ("left_wrist", "color_image", "left_wrist_image"), + ("left_wrist", "camera_info", "left_wrist_camera_info"), + ("left_wrist", "tf", "left_wrist_tf"), + ("right_wrist", "color_image", "right_wrist_image"), + ("right_wrist", "camera_info", "right_wrist_camera_info"), + ("right_wrist", "tf", "right_wrist_tf"), + ] +) diff --git a/dimos/robot/manipulators/dual_openyam/blueprints/test_learning.py b/dimos/robot/manipulators/dual_openyam/blueprints/test_learning.py index 595e8dd161..21d05ef621 100644 --- a/dimos/robot/manipulators/dual_openyam/blueprints/test_learning.py +++ b/dimos/robot/manipulators/dual_openyam/blueprints/test_learning.py @@ -13,76 +13,37 @@ # limitations under the License. -import pytest - -from dimos.experimental.memory.rust_recorder import RustRecorder +from dimos.core.coordination.blueprint_config.parser import BlueprintConfigParser from dimos.hardware.sensors.camera.module import CameraModule from dimos.robot.manipulators.dual_openyam.blueprints.learning_collection import ( - build_dual_openyam_quest_collection, + dual_openyam_quest_collection, ) -from dimos.robot.manipulators.dual_openyam.learning import DUAL_OPENYAM_COLLECTION -def test_dual_collection_declares_two_distinct_cameras_and_both_buses(tmp_path): - blueprint = build_dual_openyam_quest_collection( - recording=tmp_path / "dual.mcap", - task="fold towel", - cameras={"left_wrist_image": 0, "right_wrist_image": 1}, - left_can_port="follower_l", - right_can_port="follower_r", +def test_dual_collection_configures_both_cameras_and_buses_through_run(tmp_path): + blueprint = dual_openyam_quest_collection + parsed = BlueprintConfigParser(blueprint).parse( + [ + "--recorder.recording", + str(tmp_path / "dual"), + "--episodes.task", + "fold towel", + "--controlcoordinator.left-can-port", + "follower_l", + "--controlcoordinator.right-can-port", + "follower_r", + "--left-wrist.hardware.camera-index", + "/dev/video2", + "--right-wrist.hardware.camera-index", + "/dev/video4", + ], + environ={}, ) cameras = [atom for atom in blueprint.active_blueprints if atom.module is CameraModule] - assert issubclass(blueprint.active_blueprints[0].module, RustRecorder) - assert [camera.kwargs["hardware"].camera_index for camera in cameras] == [0, 1] - assert ( - blueprint.remapping_map[("CollectionCamera_left_wrist_image", "color_image")] - == "left_wrist_image" - ) - assert ( - blueprint.remapping_map[("CollectionCamera_right_wrist_image", "color_image")] - == "right_wrist_image" - ) - coordinator = next( - atom for atom in blueprint.active_blueprints if atom.name == "ControlCoordinator" - ) - assert coordinator.kwargs["left_can_port"] == "follower_l" - assert coordinator.kwargs["right_can_port"] == "follower_r" - - -def test_custom_profile_adds_overhead_camera_without_a_new_recorder(tmp_path): - profile = DUAL_OPENYAM_COLLECTION.model_copy(deep=True) - profile.observations["overhead"] = profile.observations[ - "observation.images.left_wrist" - ].model_copy( - update={"stream": "overhead_image"}, - ) - blueprint = build_dual_openyam_quest_collection( - profile=profile, - recording=tmp_path / "three.mcap", - task="fold towel", - cameras={"left_wrist_image": 0, "right_wrist_image": 1, "overhead_image": 2}, - ) - recorder = blueprint.active_blueprints[0] - assert {s.name for s in recorder.streams} >= { - "left_wrist_image", - "right_wrist_image", - "overhead_image", - "status", - } - assert len([atom for atom in blueprint.active_blueprints if atom.module is CameraModule]) == 3 - - -@pytest.mark.parametrize( - ("devices", "error"), - [ - ({"left_wrist_image": 0}, "missing cameras.*right_wrist_image"), - ({"left_wrist_image": 0, "right_wrist_image": 1, "typo": 2}, "unknown cameras.*typo"), - ], -) -def test_invalid_camera_bindings_fail_before_hardware_import(tmp_path, devices, error): - with pytest.raises(ValueError, match=error): - build_dual_openyam_quest_collection( - recording=tmp_path / "dual.mcap", - task="fold", - cameras=devices, - ) + assert [camera.name for camera in cameras] == ["left_wrist", "right_wrist"] + assert blueprint.remapping_map[("left_wrist", "color_image")] == "left_wrist_image" + assert blueprint.remapping_map[("right_wrist", "color_image")] == "right_wrist_image" + assert parsed.module_configs["ControlCoordinator"]["left_can_port"] == "follower_l" + assert parsed.module_configs["ControlCoordinator"]["right_can_port"] == "follower_r" + assert parsed.module_configs["left_wrist"]["hardware"]["camera_index"] == "/dev/video2" + assert parsed.module_configs["right_wrist"]["hardware"]["camera_index"] == "/dev/video4" From 10ef0d53357909cd547ad54e34bd0e90510e2a7e Mon Sep 17 00:00:00 2001 From: cc Date: Thu, 10 Sep 2026 20:15:47 -0700 Subject: [PATCH 4/4] refactor(openyam): reuse the existing dual teleop blueprint --- .../dual_openyam/blueprints/teleop.py | 86 ++++++++----------- 1 file changed, 37 insertions(+), 49 deletions(-) diff --git a/dimos/robot/manipulators/dual_openyam/blueprints/teleop.py b/dimos/robot/manipulators/dual_openyam/blueprints/teleop.py index 1df57ff9fc..47727c4d07 100644 --- a/dimos/robot/manipulators/dual_openyam/blueprints/teleop.py +++ b/dimos/robot/manipulators/dual_openyam/blueprints/teleop.py @@ -15,7 +15,7 @@ """Coupled Quest teleoperation for the complete Dual OpenYAM entity.""" from dimos.control.coordinator import TaskConfig -from dimos.core.coordination.blueprints import Blueprint, autoconnect +from dimos.core.coordination.blueprints import autoconnect from dimos.manipulation.manipulation_module import ManipulationModule from dimos.manipulation.planning.kinematics.config import PinkKinematicsConfig from dimos.robot.manipulators.common.blueprints import teleop_ik_task @@ -73,51 +73,39 @@ }, ) - -def build_dual_openyam_quest_teleop( - *, - left_can_port: str | None = None, - right_can_port: str | None = None, -) -> Blueprint: - """Build Quest teleop against mock or explicitly selected dual-CAN hardware.""" - return autoconnect( - ArmTeleopModule.blueprint(), - DualOpenYamCoordinator.blueprint( - instance_name="ControlCoordinator", - left_can_port=left_can_port, - right_can_port=right_can_port, - tasks=[ - _dual_openyam_quest_task, - TaskConfig( - name="left_arm_gripper", - type="gripper", - joint_names=[DUAL_OPENYAM_GRIPPER_JOINTS[0]], - priority=20, - stream_bind={"gripper_command": "left_gripper_command"}, - ), - TaskConfig( - name="right_arm_gripper", - type="gripper", - joint_names=[DUAL_OPENYAM_GRIPPER_JOINTS[1]], - priority=20, - stream_bind={"gripper_command": "right_gripper_command"}, - ), - dual_openyam_trajectory_task(priority=20), - ], - ), - ManipulationModule.blueprint( - model=_dual_openyam_quest_model, - kinematics=_dual_openyam_quest_pink, - visualization={"backend": "viser"}, - ), - ).remappings( - [ - (ArmTeleopModule, "left_controller_output", "left_cartesian_command"), - (ArmTeleopModule, "left_gripper_command", "left_gripper_command"), - (ArmTeleopModule, "right_controller_output", "right_cartesian_command"), - (ArmTeleopModule, "right_gripper_command", "right_gripper_command"), - ] - ) - - -teleop_quest_dual_openyam = autoconnect(build_dual_openyam_quest_teleop()) +teleop_quest_dual_openyam = autoconnect( + ArmTeleopModule.blueprint(), + DualOpenYamCoordinator.blueprint( + instance_name="ControlCoordinator", + tasks=[ + _dual_openyam_quest_task, + TaskConfig( + name="left_arm_gripper", + type="gripper", + joint_names=[DUAL_OPENYAM_GRIPPER_JOINTS[0]], + priority=20, + stream_bind={"gripper_command": "left_gripper_command"}, + ), + TaskConfig( + name="right_arm_gripper", + type="gripper", + joint_names=[DUAL_OPENYAM_GRIPPER_JOINTS[1]], + priority=20, + stream_bind={"gripper_command": "right_gripper_command"}, + ), + dual_openyam_trajectory_task(priority=20), + ], + ), + ManipulationModule.blueprint( + model=_dual_openyam_quest_model, + kinematics=_dual_openyam_quest_pink, + visualization={"backend": "viser"}, + ), +).remappings( + [ + (ArmTeleopModule, "left_controller_output", "left_cartesian_command"), + (ArmTeleopModule, "left_gripper_command", "left_gripper_command"), + (ArmTeleopModule, "right_controller_output", "right_cartesian_command"), + (ArmTeleopModule, "right_gripper_command", "right_gripper_command"), + ] +)