diff --git a/egomimic/algo/eval_determinism.py b/egomimic/algo/eval_determinism.py new file mode 100644 index 000000000..e4849a6d3 --- /dev/null +++ b/egomimic/algo/eval_determinism.py @@ -0,0 +1,109 @@ +"""Deterministic evaluation shared by the flow-matching algos (HPT, QwenVLA). + +Every RNG used while an algo is in eval mode is seeded from +``EVAL_BASE_SEED + rank * EVAL_RANK_STRIDE + pass_counter``, where the pass +counter counts ``forward_eval`` calls since the last +``reset_eval_pass_counter()`` (called from ``ModelWrapper.on_validation_start``). +Two validation passes over the same batches therefore replay the same prompts +and the same sampling noise, while different ranks and different batches +within a pass still see different draws. + +The mixin reads ``self.nets`` (an ``nn.ModuleDict``, for ``.training``), +``self.device``, ``self.annotation_key``, ``self.annotation_sampling_mode`` +and ``self.default_prompt`` from the algo it is mixed into. +""" + +from __future__ import annotations + +import os +import random + +import torch + +EVAL_BASE_SEED = 0 +EVAL_RANK_STRIDE = 1_000_003 + + +class DeterministicEvalMixin: + annotation_key: str | None = None + annotation_sampling_mode: str = "random" + default_prompt: str = "" + + def reset_eval_pass_counter(self) -> None: + """Start a new validation pass: the next ``forward_eval`` call is pass 0. + + Called from ``ModelWrapper.on_validation_start``. Two passes over the + same loader (same order, same number of batches) then draw the same + prompts and the same sampling noise. + """ + self._eval_pass_counter = 0 + + @staticmethod + def _eval_rank() -> int: + """Distributed rank, or 0 outside a process group.""" + if torch.distributed.is_available() and torch.distributed.is_initialized(): + return torch.distributed.get_rank() + for var in ("RANK", "LOCAL_RANK", "SLURM_PROCID"): + value = os.environ.get(var) + if value is not None and value.lstrip("-").isdigit(): + return int(value) + return 0 + + def _eval_pass_seed(self) -> int: + """Seed for the current eval pass: base + rank stride + pass counter.""" + return ( + int(getattr(self, "eval_base_seed", EVAL_BASE_SEED)) + + self._eval_rank() * EVAL_RANK_STRIDE + + int(getattr(self, "_eval_pass_counter", 0)) + ) + + def _eval_device(self) -> torch.device: + return ( + torch.device(self.device) + if self.device is not None + else torch.device("cpu") + ) + + def _eval_generator(self) -> torch.Generator: + """Seeded generator on the device the sampling noise is drawn on.""" + device = self._eval_device() + generator = torch.Generator(device=device) + generator.manual_seed(self._eval_pass_seed()) + return generator + + def _eval_prompt_rng(self) -> random.Random: + """Seeded Python RNG for annotation choice during eval.""" + return random.Random(self._eval_pass_seed()) + + def _eval_fork_devices(self) -> list: + """Devices ``torch.random.fork_rng`` must save/restore (CPU is implicit).""" + device = self._eval_device() + if device.type != "cuda": + return [] + return [ + device.index if device.index is not None else torch.cuda.current_device() + ] + + def _build_prompts(self, _batch, batch_size: int) -> list[str]: + """Sample one annotation per batch item, falling back to default_prompt + on empty / missing annotations. Mirrors the Pi algo flow. + + In eval the choice comes from a per-pass seeded RNG regardless of + ``annotation_sampling_mode``: that keeps prompt variety (unlike always + taking ``sample[0]``) while making two validation passes over the same + batches produce identical prompt lists. Training is unchanged. + """ + if self.annotation_key is None or self.annotation_key not in _batch: + return [self.default_prompt] * batch_size + eval_rng = None if self.nets.training else self._eval_prompt_rng() + prompts = [] + for sample in _batch[self.annotation_key]: + if not sample: + prompts.append(self.default_prompt) + elif eval_rng is not None: + prompts.append(sample[eval_rng.randint(0, len(sample) - 1)]) + elif self.annotation_sampling_mode == "random": + prompts.append(sample[random.randint(0, len(sample) - 1)]) + else: # "first" + prompts.append(sample[0]) + return prompts diff --git a/egomimic/algo/hpt.py b/egomimic/algo/hpt.py index 994e21723..e295e59db 100644 --- a/egomimic/algo/hpt.py +++ b/egomimic/algo/hpt.py @@ -1,5 +1,4 @@ import os -import random from collections import OrderedDict from functools import partial from typing import Literal, Union @@ -14,6 +13,7 @@ from tslearn.metrics import SoftDTWLossPyTorch from egomimic.algo.algo import Algo +from egomimic.algo.eval_determinism import DeterministicEvalMixin from egomimic.models.hpt_nets import ( MultiheadAttention, SimpleTransformer, @@ -29,16 +29,6 @@ # it has exactly one consumer, below). STD_SCALE = 0.02 -# Deterministic evaluation. Every RNG used while the algo is in eval mode is -# seeded from ``EVAL_BASE_SEED + rank * EVAL_RANK_STRIDE + pass_counter``, where -# the pass counter counts ``forward_eval`` calls since the last -# ``reset_eval_pass_counter()`` (called from ``ModelWrapper.on_validation_start``). -# Two validation passes over the same batches therefore replay the same prompts -# and the same sampling noise, while different ranks and different batches -# within a pass still see different draws. -EVAL_BASE_SEED = 0 -EVAL_RANK_STRIDE = 1_000_003 - class HPTModel(nn.Module): """ @@ -831,7 +821,7 @@ def load_pretrained(self, checkpoint_path): self.load_trunk(os.path.join(checkpoint_path, "trunk.pth")) -class HPT(Algo): +class HPT(DeterministicEvalMixin, Algo): """ """ @property @@ -1046,88 +1036,6 @@ def __init__( self.training_step = 0 - # ------------------------------------------------------------------ - # Deterministic evaluation - # ------------------------------------------------------------------ - def reset_eval_pass_counter(self) -> None: - """Start a new validation pass: the next ``forward_eval`` call is pass 0. - - Called from ``ModelWrapper.on_validation_start``. Two passes over the - same loader (same order, same number of batches) then draw the same - prompts and the same sampling noise. - """ - self._eval_pass_counter = 0 - - @staticmethod - def _eval_rank() -> int: - """Distributed rank, or 0 outside a process group.""" - if torch.distributed.is_available() and torch.distributed.is_initialized(): - return torch.distributed.get_rank() - for var in ("RANK", "LOCAL_RANK", "SLURM_PROCID"): - value = os.environ.get(var) - if value is not None and value.lstrip("-").isdigit(): - return int(value) - return 0 - - def _eval_pass_seed(self) -> int: - """Seed for the current eval pass: base + rank stride + pass counter.""" - return ( - int(getattr(self, "eval_base_seed", EVAL_BASE_SEED)) - + self._eval_rank() * EVAL_RANK_STRIDE - + int(getattr(self, "_eval_pass_counter", 0)) - ) - - def _eval_device(self) -> torch.device: - return ( - torch.device(self.device) - if self.device is not None - else torch.device("cpu") - ) - - def _eval_generator(self) -> torch.Generator: - """Seeded generator on the device the sampling noise is drawn on.""" - device = self._eval_device() - generator = torch.Generator(device=device) - generator.manual_seed(self._eval_pass_seed()) - return generator - - def _eval_prompt_rng(self) -> random.Random: - """Seeded Python RNG for annotation choice during eval.""" - return random.Random(self._eval_pass_seed()) - - def _eval_fork_devices(self) -> list: - """Devices ``torch.random.fork_rng`` must save/restore (CPU is implicit).""" - device = self._eval_device() - if device.type != "cuda": - return [] - return [ - device.index if device.index is not None else torch.cuda.current_device() - ] - - def _build_prompts(self, _batch, batch_size: int) -> list[str]: - """Sample one annotation per batch item, falling back to default_prompt - on empty / missing annotations. Mirrors the Pi algo flow. - - In eval the choice comes from a per-pass seeded RNG regardless of - ``annotation_sampling_mode``: that keeps prompt variety (unlike always - taking ``sample[0]``) while making two validation passes over the same - batches produce identical prompt lists. Training is unchanged. - """ - if self.annotation_key is None or self.annotation_key not in _batch: - return [self.default_prompt] * batch_size - eval_rng = None if self.nets.training else self._eval_prompt_rng() - prompts = [] - for sample in _batch[self.annotation_key]: - if not sample: - prompts.append(self.default_prompt) - elif eval_rng is not None: - prompts.append(sample[eval_rng.randint(0, len(sample) - 1)]) - elif self.annotation_sampling_mode == "random": - prompts.append(sample[random.randint(0, len(sample) - 1)]) - else: # "first" - prompts.append(sample[0]) - return prompts - @override def process_batch_for_training(self, batch): """ diff --git a/egomimic/algo/qwenvla.py b/egomimic/algo/qwenvla.py new file mode 100644 index 000000000..b92be3c90 --- /dev/null +++ b/egomimic/algo/qwenvla.py @@ -0,0 +1,331 @@ +"""QwenVLA: Qwen 3.5 VLM + layer-wise cross-DiT flow head (starVLA QwenPI port). + +Design: docs/hpt-experiments/2026-09-16_qwenvla_design.md. Modeled on the PI +algo (prompt assembly with an embodiment block, one loop per embodiment, +norm stats, the unnormalised-prediction contract ``eval_hpt.py`` reads) and +fed HPT-style batches (native action layout, dotted dataset keys, proprio +``(B, D)`` or ``(B, K, D)`` from the ``proprio_history`` keymap). +""" + +from __future__ import annotations + +import logging +from collections import OrderedDict +from typing import Literal, Optional + +import torch +import torch.nn as nn +from overrides import override + +from egomimic.algo.algo import Algo +from egomimic.algo.eval_determinism import DeterministicEvalMixin +from egomimic.models.hpt_nets import verify_pretrained_weights +from egomimic.models.image_augs import PerSampleAugs +from egomimic.models.qwenvla_nets import QwenVLAModel, QwenVLBackbone +from egomimic.rldb.embodiment.embodiment import get_embodiment, get_embodiment_id +from egomimic.utils.action_utils import pad_to_width + +logger = logging.getLogger(__name__) + + +class QwenVLA(DeterministicEvalMixin, Algo): + @property + def device(self): + return self._device + + @device.setter + def device(self, value): + """Move the algo (and its nets) to ``value``; written by ``ModelWrapper`` + once Lightning knows the rank's device (see ``__init__``).""" + self._device = torch.device(value) if value is not None else torch.device("cpu") + nets = getattr(self, "nets", None) + if nets is not None: + nets.to(self._device) + + def __init__( + self, + norm_stats, + backbone: QwenVLBackbone, + head, # functools.partial of LayerwiseFMHead (hydra ``_partial_: true``) + domains: list, + dims: dict, + ac_keys: dict, + action_width: Optional[int] = None, + train_image_augs=None, + eval_image_augs=None, + annotation_key: Optional[str] = None, + annotation_sampling_mode: Literal["random", "first"] = "random", + default_prompt: str = "", + embodiment_label: bool = True, + control_mode: Optional[dict] = None, + use_pad_mask: bool = False, + history_len: int = 1, + history_dropout: float = 0.0, + **kwargs, + ): + self.nets = nn.ModuleDict() + self.norm_stats = norm_stats + self.domains = list(domains) + self.dims = {d: dict(dims[d]) for d in self.domains} + widths = {d: int(self.dims[d]["action"]) for d in self.domains} + self.action_width = ( + int(action_width) if action_width is not None else max(widths.values()) + ) + for d, w in widths.items(): + if w > self.action_width: + raise ValueError( + f"{d} action is {w}-wide but action_width={self.action_width}" + ) + self.annotation_key = annotation_key + self.annotation_sampling_mode = annotation_sampling_mode + self.default_prompt = default_prompt + self.embodiment_label = bool(embodiment_label) + self.control_mode = dict(control_mode) if control_mode else None + self.use_pad_mask = bool(use_pad_mask) + self.train_image_augs = ( + PerSampleAugs(train_image_augs) if train_image_augs is not None else None + ) + self.eval_image_augs = eval_image_augs + self.is_6dof = kwargs.get("6dof", False) + # Read by eval_hpt.py; QwenVLA has one head, no auxiliary / shared keys. + self.shared_ac_key = None + self.auxiliary_ac_keys = {} + # Construct on the CPU. Lightning learns each rank's GPU only after the + # model exists and writes it to ``device`` (``ModelWrapper.on_fit_start`` + # / ``on_validation_start``), whose setter moves the nets. Picking + # ``cuda`` here would put every DDP rank's 4.8 GB of fp32 VLM weights on + # cuda:0 until then, which overflows it on a checkpoint resume. + self._device = torch.device(kwargs.get("device", "cpu")) + + state_dims = { + d: int(self.dims[d]["proprio"]) + for d in self.domains + if self.dims[d].get("proprio") + } + head_module = head( + action_width=self.action_width, + vlm_hidden=backbone.hidden_size, + num_layers=backbone.num_layers, + state_dims=state_dims or None, + history_len=history_len, + history_dropout=history_dropout, + ) + self.nets["policy"] = QwenVLAModel(backbone, head_module) + self.nets = self.nets.float().to(self._device) + verify_pretrained_weights(self.nets["policy"]) + + self.ac_keys = dict(ac_keys) + self.camera_keys, self.proprio_keys, self.lang_keys = {}, {}, {} + for embodiment in self.domains: + embodiment_id = get_embodiment_id(embodiment) + self.camera_keys[embodiment_id] = [] + self.proprio_keys[embodiment_id] = [] + self.lang_keys[embodiment_id] = [] + for key in norm_stats.keys_of_type("action_keys", embodiment_id): + if ( + norm_stats.is_key_with_embodiment(key, embodiment_id) + and key == self.ac_keys[embodiment] + ): + self.ac_keys[embodiment_id] = key + for key in norm_stats.keys_of_type("camera_keys", embodiment_id): + if norm_stats.is_key_with_embodiment(key, embodiment_id): + self.camera_keys[embodiment_id].append(key) + for key in norm_stats.keys_of_type("proprio_keys", embodiment_id): + if norm_stats.is_key_with_embodiment(key, embodiment_id): + self.proprio_keys[embodiment_id].append(key) + for key in norm_stats.keys_of_type("lang_keys", embodiment_id): + if norm_stats.is_key_with_embodiment(key, embodiment_id): + self.lang_keys[embodiment_id].append(key) + if embodiment_id not in self.ac_keys: + raise KeyError( + f"ac_keys[{embodiment}]={self.ac_keys[embodiment]!r} is not an action " + "key of the norm stats" + ) + self.training_step = 0 + + # ------------------------------------------------------------------ + # prompts + # ------------------------------------------------------------------ + def _control_mode_for(self, embodiment_name: str) -> str: + for key, value in self.control_mode.items(): + if key.lower() in embodiment_name.lower(): + return value + raise ValueError( + f"control_mode has no entry matching embodiment {embodiment_name!r} " + f"(declared keys: {sorted(self.control_mode)})" + ) + + def compose_prompt(self, embodiment_name: str, annotation: str) -> str: + """``Embodiment: human bimanual. [Control mode: X.] Task: ``; + the task block is dropped when the annotation is empty.""" + blocks = [] + if self.embodiment_label: + blocks.append(f"Embodiment: {embodiment_name.lower().replace('_', ' ')}.") + if self.control_mode: + blocks.append(f"Control mode: {self._control_mode_for(embodiment_name)}.") + annotation = (annotation or "").strip() + if annotation: + blocks.append(f"Task: {annotation}") + return " ".join(blocks) + + # ------------------------------------------------------------------ + # batches + # ------------------------------------------------------------------ + @override + def process_batch_for_training(self, batch): + processed_batch = {} + for embodiment_name, _batch in batch.items(): + embodiment_id = get_embodiment_id(embodiment_name) + out = {} + for key, value in _batch.items(): + key_name = ( + self.norm_stats.zarr_key_to_keyname(key, embodiment_id) or key + ) + out[key_name] = value + ac_key = self.ac_keys[embodiment_id] + if ac_key not in out or out[ac_key].ndim != 3: + raise ValueError(f"expected a (B, T, D) action under {ac_key!r}") + B, S, _ = out[ac_key].shape + action_pad_mask = out.pop("action_pad_mask", None) + out.pop("proprio_history_mask", None) # front padding == history dropout + if self.use_pad_mask and action_pad_mask is not None: + if action_pad_mask.shape[-1] != S: + raise ValueError( + f"action_pad_mask length {action_pad_mask.shape[-1]} != horizon {S}" + ) + out["pad_mask"] = action_pad_mask[..., None] + else: + out["pad_mask"] = torch.ones(B, S, 1) + out["sampled_prompt"] = self._build_prompts(_batch, B) + out["embodiment"] = torch.tensor([embodiment_id], dtype=torch.int64) + for key, value in out.items(): + if isinstance(value, torch.Tensor): + value = value.to(self.device) + if value.is_floating_point(): + value = value.float() + out[key] = value + processed_batch[embodiment_id] = out + return processed_batch + + def _apply_image_augs(self, images: torch.Tensor) -> torch.Tensor: + if self.nets.training and self.train_image_augs is not None: + return self.train_image_augs(images) + if not self.nets.training and self.eval_image_augs is not None: + return self.eval_image_augs(images) + return images + + def _to_model_data(self, _batch: dict, embodiment_id: int) -> dict: + """Processed batch -> the dict ``QwenVLAModel`` consumes.""" + name = get_embodiment(embodiment_id).lower() + frames = [ + self._apply_image_augs(_batch[key]) + for key in self.camera_keys[embodiment_id] + if key in _batch + ] + if not frames: + raise KeyError( + f"no camera of {self.camera_keys[embodiment_id]} in the batch keys {sorted(_batch)}" + ) + images = torch.stack(frames, dim=1) # (B, F, 3, H, W) + + parts = [] + for key in self.proprio_keys[embodiment_id]: + if key not in _batch: + continue + value = _batch[key] + parts.append(value[:, None] if value.ndim == 2 else value) # (B, K, D) + state = None + if parts: + steps = {p.shape[1] for p in parts} + if len(steps) != 1: + raise ValueError( + f"proprio keys carry different history lengths: {steps}" + ) + state = torch.cat(parts, dim=-1) + expected = int(self.dims[name]["proprio"]) + if state.shape[-1] != expected: + raise ValueError( + f"{name} proprio is {state.shape[-1]}-wide, dims says {expected}" + ) + + ac_key = self.ac_keys[embodiment_id] + action = _batch[ac_key] + native_width = action.shape[-1] + action = pad_to_width(action, self.action_width) + width_mask = torch.zeros(1, 1, self.action_width, device=action.device) + width_mask[..., :native_width] = 1.0 + loss_mask = _batch["pad_mask"].to(action.device) * width_mask + + prompts = [self.compose_prompt(name, p) for p in _batch["sampled_prompt"]] + return { + "images": images, + "prompts": prompts, + "state": state, + "action": action, + "loss_mask": loss_mask, + "embodiment_name": name, + } + + # ------------------------------------------------------------------ + # train / eval + # ------------------------------------------------------------------ + @override + def forward_training(self, batch): + predictions = OrderedDict() + self.training_step += 1 + for embodiment_id, _batch in batch.items(): + name = get_embodiment(embodiment_id).lower() + ac_key = self.ac_keys[embodiment_id] + data = self._to_model_data(_batch, embodiment_id) + predictions[f"{name}_{ac_key}"] = _batch[ac_key] + predictions[f"{name}_loss"] = self.nets["policy"].compute_loss(data) + return predictions + + @override + def forward_eval(self, batch): + unnorm_preds = {} + generator = self._eval_generator() + pass_seed = self._eval_pass_seed() + fork_devices = self._eval_fork_devices() + with torch.no_grad(): + for embodiment_id, _batch in batch.items(): + name = get_embodiment(embodiment_id).lower() + ac_key = self.ac_keys[embodiment_id] + data = self._to_model_data(_batch, embodiment_id) + # Val loss draws noise / t from the global RNG inside the head: + # fork + seed it so the reported loss is reproducible. + with torch.random.fork_rng(devices=fork_devices): + torch.manual_seed(pass_seed + 1) + unnorm_preds[f"{name}_loss"] = self.nets["policy"].compute_loss( + data + ) + pred = self.nets["policy"].sample(data, generator=generator) + ref = _batch[ac_key] + _, T, D = ref.shape + unnorm = self.norm_stats.unnormalize( + {ac_key: pred[:, :T, :D]}, embodiment_id + ) + for key, value in unnorm.items(): + unnorm_preds[f"{name}_{key}"] = value + self._eval_pass_counter = int(getattr(self, "_eval_pass_counter", 0)) + 1 + return unnorm_preds + + @override + def compute_losses(self, predictions, batch): + loss_dict = OrderedDict() + total = torch.tensor(0.0, device=self.device) + for embodiment_id in batch: + name = get_embodiment(embodiment_id).lower() + bc_loss = predictions[f"{name}_loss"] + total = total + bc_loss + loss_dict[f"{name}_loss"] = bc_loss + loss_dict["action_loss"] = total / len(self.domains) + return loss_dict + + @override + def log_info(self, info): + log = OrderedDict() + log["Loss"] = info["losses"]["action_loss"].item() + for loss_key, loss in info["losses"].items(): + log[loss_key] = loss.item() + return log diff --git a/egomimic/eval/eval_video.py b/egomimic/eval/eval_video.py index 9fbcc6172..49042f9a9 100644 --- a/egomimic/eval/eval_video.py +++ b/egomimic/eval/eval_video.py @@ -87,15 +87,18 @@ def compute_metrics_and_viz(self, batch, do_viz=True): raise NotImplementedError def _video_fps(self, source_fps: int = 30) -> int: - """Playback fps compensating for the DistributedSampler stride. - - Distributed validation deals the (unshuffled) val indices round-robin, - so rank 0's consecutive rendered frames are ``world_size`` source - frames apart. Writing them at the raw source fps produces a - world_size-times timelapse; scale playback down to restore wall-clock - pacing.""" - world = max(1, int(getattr(self.trainer, "world_size", 1) or 1)) - return max(1, round(source_fps / world)) + """Playback fps of the overlay video: the source rate. + + The val heads come back from ``val_dataloader()`` inside + ``CombinedLoader``s, which Lightning does NOT wrap in a + ``DistributedSampler`` (every rank runs every val batch; measured + 2026-09-16: 1 rank = 223 val steps, 8 ranks = 8 x 223), so rank 0 + renders every source frame and the video plays in real time at the + source rate. An earlier version divided by ``world_size`` to undo a + sampler stride that never applied, which made every multi-GPU video a + ``world_size``-times slow-motion (416 frames at 4 fps instead of 30). + """ + return int(source_fps) def _write_video(self, key, frames) -> None: path = os.path.join( diff --git a/egomimic/hydra_configs/data/mecka_fold_flagship_opsplit_hpt_keypoints.yaml b/egomimic/hydra_configs/data/mecka_fold_flagship_opsplit_hpt_keypoints.yaml new file mode 100644 index 000000000..1bbed47b5 --- /dev/null +++ b/egomimic/hydra_configs/data/mecka_fold_flagship_opsplit_hpt_keypoints.yaml @@ -0,0 +1,26 @@ +# mecka_fold_flagship_opsplit_hpt_6d with the wrist-frame KEYPOINT pipeline +# (the keymap / transform of mecka_fold_freeform_opsplit_hpt_keypoints): +# actions_keypoints (100, 144) wrist-first wrist-frame keypoint +# action, per hand [wrist xyz+rot6d +# (9) | 21 keypoints (63)] +# observations.state.keypoints (144,) per hand [wrist pose in HEAD +# frame | keypoints in wrist frame] +# No `observations.state.ee_pose` (include_ee_pose off), so +# pad_proprio_gripper is off. The operator split, filters, val loaders and +# video episodes are inherited unchanged. Pairs with +# model=hpt_bc_keypoints_wrist_300M / qwenvla_bc_mecka_kp_0p8b and +# evaluator=eval_hpt_keypoints_wrist. Norm stats are data-specific: +# precompute them for this config (egomimic/scripts/precompute_norm_stats.py). +defaults: + - mecka_fold_flagship_opsplit_hpt_6d + - _self_ + +train_datasets: + human_bimanual: + resolver: + key_map: + keymap_mode: keypoints + transform_list: + mode: keypoints_wristframe_6d + fix_mecka_left_wrist: true + pad_proprio_gripper: false diff --git a/egomimic/hydra_configs/data/mecka_fold_flagship_top3_hpt_keypoints.yaml b/egomimic/hydra_configs/data/mecka_fold_flagship_top3_hpt_keypoints.yaml new file mode 100644 index 000000000..89a3ac0cb --- /dev/null +++ b/egomimic/hydra_configs/data/mecka_fold_flagship_top3_hpt_keypoints.yaml @@ -0,0 +1,17 @@ +# mecka_fold_flagship_top3_hpt_6d (the 3-operator flagship train split) with +# the wrist-frame KEYPOINT pipeline; see +# mecka_fold_flagship_opsplit_hpt_keypoints for the key layout. Norm stats: +# logs/hpt_architecture/norm_stats/mecka_fold_flagship_top3_hpt_keypoints/norm_stats +defaults: + - mecka_fold_flagship_top3_hpt_6d + - _self_ + +train_datasets: + human_bimanual: + resolver: + key_map: + keymap_mode: keypoints + transform_list: + mode: keypoints_wristframe_6d + fix_mecka_left_wrist: true + pad_proprio_gripper: false diff --git a/egomimic/hydra_configs/model/hpt_bc_keypoints_wrist_300M.yaml b/egomimic/hydra_configs/model/hpt_bc_keypoints_wrist_300M.yaml index 1b7f4922b..c992e14b6 100644 --- a/egomimic/hydra_configs/model/hpt_bc_keypoints_wrist_300M.yaml +++ b/egomimic/hydra_configs/model/hpt_bc_keypoints_wrist_300M.yaml @@ -126,6 +126,11 @@ robomimic_model: input_dim: ${model.robomimic_model.dims.human_bimanual.proprio} output_dim: 840 widths: [840] + # Proprio observation history (as hpt_bc_mecka_6d_300M's state_ee_pose + # stem): K tokens, learned zero-init time-step embedding, whole-past + # dropout. Must equal the data config's `proprio_history`; inert at 1. + history_len: 1 + history_dropout: 0.2 specs: random_horizon_masking: false cross_attn: @@ -159,6 +164,10 @@ robomimic_model: mean: [0.485, 0.456, 0.406] std: [0.229, 0.224, 0.225] +# LR multiplier for the image encoders' pretrained backbones; 1.0 keeps the +# single optimizer param group (ModelWrapper.configure_optimizers). +backbone_lr_scale: 1.0 + optimizer: _target_: torch.optim.AdamW _partial_: true diff --git a/egomimic/hydra_configs/model/qwenvla_bc_mecka_6d_0p8b.yaml b/egomimic/hydra_configs/model/qwenvla_bc_mecka_6d_0p8b.yaml new file mode 100644 index 000000000..a6006c87f --- /dev/null +++ b/egomimic/hydra_configs/model/qwenvla_bc_mecka_6d_0p8b.yaml @@ -0,0 +1,105 @@ +# QWENVLA_BC_MECKA_6D_0P8B: QwenVLA (our port of starVLA's QwenPI) BC on the +# 18-D wrist-frame 6D cartesian action of the mecka fold-clothes data. +# Design: docs/hpt-experiments/2026-09-16_qwenvla_design.md. +# +# backbone Qwen/Qwen3.5-0.8B (24 text layers, hidden 1024), FULL fine-tune +# at backbone_lr_scale x the head LR, fp32 master weights under +# Lightning's bf16 autocast, gradient checkpointing on. One chat +# turn per sample: [front frame] + "Embodiment: human bimanual. +# Task: "; 640x352 -> 220 visual tokens. +# head LayerwiseFMHead: 24 DiT blocks at the VLM width (projectors are +# the identity), block i cross-attends to VLM layer i, odd blocks +# self-attend; 32 register tokens + 100 action tokens; state = one +# MLP token per proprio step, prepended (starVLA's way, NOT text). +# Flow matching as fm_policy.py (Beta(1.5, 1) time, Euler, 10 steps). +# data the same batches as hpt_bc_mecka_6d_300M: `actions_cartesian` +# (18), `observations.state.ee_pose` (20, 6D-encoded + grip pads), +# `front_img_1` in [0, 1]; annotations sampled per item. +# +# Run (the HPT flagship recipe with this model): +# python egomimic/trainHydra.py --config-name=train_zarr_mecka_flagship_6d_qwenvla \ +# data=mecka_fold_flagship_top3_hpt_6d \ +# norm_stats.precomputed_norm_path=/norm_stats +# +# Proprio history: set BOTH +# data.train_datasets.human_bimanual.resolver.key_map.proprio_history=K +# model.robomimic_model.history_len=K +_target_: egomimic.pl_utils.pl_model.ModelWrapper +robomimic_model: + _target_: egomimic.algo.qwenvla.QwenVLA + domains: ["human_bimanual"] + dims: + human_bimanual: {proprio: 20, action: 18} + ac_keys: + human_bimanual: "actions_cartesian" + # Shared head width; every embodiment's action is zero-padded to it and the + # loss is masked to the real dims. + action_width: 18 + 6dof: true + + # Prompt: "Embodiment: . [Control mode: .] Task: ". + annotation_key: "annotations" + annotation_sampling_mode: "random" + default_prompt: "" + embodiment_label: true + control_mode: null + + # Mask the repeat-last padding at episode tails out of the loss. + use_pad_mask: false + + # State tokens in the DiT: K proprio steps (= the keymap's proprio_history), + # learned zero-init per-step embedding, whole-context history dropout. + history_len: 1 + history_dropout: 0.2 + + backbone: + _target_: egomimic.models.qwenvla_nets.QwenVLBackbone + model_name: "Qwen/Qwen3.5-0.8B" + dtype: "float32" + freeze: false + trainable_layers: 0 + gradient_checkpointing: true + image_size: [352, 640] # (H, W) -> 220 visual tokens per frame + max_text_tokens: 128 + num_layers: null # all 24 text layers -> 24 DiT blocks + + head: + _target_: egomimic.models.layerwise_dit.LayerwiseFMHead + _partial_: true # QwenVLA fills action_width / vlm_hidden / num_layers / state_dims + action_horizon: 100 + dit_hidden: null # null = VLM hidden (QwenPI); 1024 at 2B (QwenPI_v3) + head_dim: 64 + num_register_tokens: 32 + dropout: 0.2 + final_dropout: true + interleave_self_attention: true + adaln_zero: false + # 50 to match the HPT baseline head (hpt_bc_mecka_6d_300M's FMPolicy): + # val metrics are computed from SAMPLED actions, so the step count is part + # of the comparison. 10 steps is a free eval-only sweep on a checkpoint. + num_inference_steps: 50 + time_dist: "beta" + repeated_diffusion_steps: 1 + + # The Qwen processor normalizes (frames must reach it in [0, 1]), so no + # ImageNet Normalize; only the ColorJitter is ours, applied per sample. + train_image_augs: + _target_: torchvision.transforms.Compose + transforms: + - _target_: torchvision.transforms.ColorJitter + brightness: 0.1 + contrast: 0.1 + saturation: 0.1 + hue: 0.05 + eval_image_augs: null + +# VLM LR = backbone_lr_scale x optimizer.lr (starVLA: 1e-5 vs 1e-4). +backbone_lr_scale: 0.1 + +optimizer: + _target_: torch.optim.AdamW + _partial_: true + lr: 1e-4 + weight_decay: 0.0001 + +scheduler: null diff --git a/egomimic/hydra_configs/model/qwenvla_bc_mecka_6d_2b.yaml b/egomimic/hydra_configs/model/qwenvla_bc_mecka_6d_2b.yaml new file mode 100644 index 000000000..79a7b6db6 --- /dev/null +++ b/egomimic/hydra_configs/model/qwenvla_bc_mecka_6d_2b.yaml @@ -0,0 +1,17 @@ +# QWENVLA_BC_MECKA_6D_2B: the 0.8B config with Qwen/Qwen3.5-2B (hidden 2048). +# The head runs at 1024 through per-layer LayerNorm+Linear projectors +# (QwenPI_v3's compressed head, ~350M + 50M instead of a 1.4B DiT). A full +# fine-tune of 2B (~32 GB of fp32 weights + grads + Adam states) does not fit +# the lab's 48 GB A40 / L40S, so the VLM is frozen except its last 8 text +# layers (`trainable_layers`); set `freeze: false` only on an 80 GB card. +defaults: + - qwenvla_bc_mecka_6d_0p8b + - _self_ + +robomimic_model: + backbone: + model_name: "Qwen/Qwen3.5-2B" + freeze: true + trainable_layers: 8 + head: + dit_hidden: 1024 diff --git a/egomimic/hydra_configs/model/qwenvla_bc_mecka_kp_0p8b.yaml b/egomimic/hydra_configs/model/qwenvla_bc_mecka_kp_0p8b.yaml new file mode 100644 index 000000000..0743d1bb4 --- /dev/null +++ b/egomimic/hydra_configs/model/qwenvla_bc_mecka_kp_0p8b.yaml @@ -0,0 +1,18 @@ +# QWENVLA_BC_MECKA_KP_0P8B: qwenvla_bc_mecka_6d_0p8b on the 144-D wrist-frame +# MANO keypoint action (the layout of hpt_bc_keypoints_wrist_300M): per hand +# [wrist xyz+rot6d in its own wrist frame (9) | 21 keypoints in that wrist +# frame (63)] x {left, right}. Proprio is the 144-D +# `observations.state.keypoints` (wrist pose in HEAD frame + keypoints in +# wrist frame, per hand), one state token per proprio step in the DiT. +# Pair with data=mecka_fold_flagship_{opsplit,top3}_hpt_keypoints and +# evaluator=eval_hpt_keypoints_wrist (recipe: train_zarr_mecka_flagship_kp_qwenvla). +defaults: + - qwenvla_bc_mecka_6d_0p8b + - _self_ +robomimic_model: + dims: + human_bimanual: {proprio: 144, action: 144} + ac_keys: + human_bimanual: "actions_keypoints" + action_width: 144 + 6dof: false diff --git a/egomimic/hydra_configs/model/qwenvla_cotrain_eva_mecka_6d.yaml b/egomimic/hydra_configs/model/qwenvla_cotrain_eva_mecka_6d.yaml new file mode 100644 index 000000000..ae43e0640 --- /dev/null +++ b/egomimic/hydra_configs/model/qwenvla_cotrain_eva_mecka_6d.yaml @@ -0,0 +1,19 @@ +# QWENVLA_COTRAIN_EVA_MECKA_6D: eva + mecka cotrain on the cartesian 6D action +# (eva 20-D with grippers, human 18-D), one shared 20-wide head, loss masked +# to each embodiment's real dims; the prompt names the embodiment. Pair with +# the data config the pi0.5 cotrain uses (see pi0.5_cotrain_eva_aria_6d.yaml +# for the transform-mode overrides) and evaluator=eval_hpt_wristframe_6d. +# Eva's three cameras all go into the chat turn (3 x 220 visual tokens). +defaults: + - qwenvla_bc_mecka_6d_0p8b + - _self_ + +robomimic_model: + domains: ["eva_bimanual", "human_bimanual"] + dims: + eva_bimanual: {proprio: 20, action: 20} + human_bimanual: {proprio: 20, action: 18} + ac_keys: + eva_bimanual: "actions_cartesian" + human_bimanual: "actions_cartesian" + action_width: 20 diff --git a/egomimic/hydra_configs/train_zarr_mecka_flagship_6d_qwenvla.yaml b/egomimic/hydra_configs/train_zarr_mecka_flagship_6d_qwenvla.yaml new file mode 100644 index 000000000..8f1859ef3 --- /dev/null +++ b/egomimic/hydra_configs/train_zarr_mecka_flagship_6d_qwenvla.yaml @@ -0,0 +1,19 @@ +# Composed recipe: QwenVLA (Qwen 3.5-0.8B + layer-wise cross-DiT) BC on the +# 18-D wrist-frame 6D-rotation cartesian action, mecka FLAGSHIP fold-clothes, +# seen/unseen operator split. Same data, evaluators and val heads as +# train_zarr_mecka_flagship_6d_hpt, so runs compare directly with the HPT +# baseline (docs/hpt-experiments/2026-09-15_hpt_plan.md, "Baseline run"). +# +# python egomimic/trainHydra.py --config-name=train_zarr_mecka_flagship_6d_qwenvla \ +# data=mecka_fold_flagship_top3_hpt_6d \ +# norm_stats.precomputed_norm_path=/norm_stats +# +# Norm stats are data-only: reuse the baseline's +# logs/hpt_architecture/norm_stats/mecka_fold_flagship_top3_hpt_6d/norm_stats. +defaults: + - train_zarr_cartesian + - evaluator@train_viz_evaluator: train_viz_hpt_wristframe_6d + - override model: qwenvla_bc_mecka_6d_0p8b + - override data: mecka_fold_flagship_opsplit_hpt_6d + - override evaluator: eval_hpt_wristframe_6d + - _self_ diff --git a/egomimic/hydra_configs/train_zarr_mecka_flagship_kp_hpt.yaml b/egomimic/hydra_configs/train_zarr_mecka_flagship_kp_hpt.yaml new file mode 100644 index 000000000..251dc0b42 --- /dev/null +++ b/egomimic/hydra_configs/train_zarr_mecka_flagship_kp_hpt.yaml @@ -0,0 +1,23 @@ +# Composed recipe: HPT 300M flow-matching BC on the 144-D wrist-frame MANO +# keypoint action, mecka FLAGSHIP fold-clothes, seen/unseen operator split. +# The keypoint counterpart of train_zarr_mecka_flagship_6d_hpt and the HPT +# reference arm of the QwenVLA comparison (train_zarr_mecka_flagship_kp_qwenvla +# is the same recipe with the model swapped). +# +# python egomimic/trainHydra.py --config-name=train_zarr_mecka_flagship_kp_hpt \ +# data=mecka_fold_flagship_top3_hpt_keypoints \ +# norm_stats.precomputed_norm_path=/norm_stats +# +# The flagship annotations are empty, so every prompt is the default one +# ("fold clothes"); the QwenVLA recipe uses the same string. +defaults: + - train_zarr_cartesian + - evaluator@train_viz_evaluator: train_viz_hpt_keypoints_wrist + - override model: hpt_bc_keypoints_wrist_300M + - override data: mecka_fold_flagship_opsplit_hpt_keypoints + - override evaluator: eval_hpt_keypoints_wrist + - _self_ + +model: + robomimic_model: + default_prompt: "fold clothes" diff --git a/egomimic/hydra_configs/train_zarr_mecka_flagship_kp_qwenvla.yaml b/egomimic/hydra_configs/train_zarr_mecka_flagship_kp_qwenvla.yaml new file mode 100644 index 000000000..1fcc6dd51 --- /dev/null +++ b/egomimic/hydra_configs/train_zarr_mecka_flagship_kp_qwenvla.yaml @@ -0,0 +1,23 @@ +# Composed recipe: QwenVLA (Qwen 3.5-0.8B + layer-wise cross-DiT) BC on the +# 144-D wrist-frame MANO keypoint action, mecka FLAGSHIP fold-clothes, +# seen/unseen operator split. Same data, evaluators and val heads as +# train_zarr_mecka_flagship_kp_hpt, so runs compare directly with that HPT arm. +# +# python egomimic/trainHydra.py --config-name=train_zarr_mecka_flagship_kp_qwenvla \ +# data=mecka_fold_flagship_top3_hpt_keypoints \ +# norm_stats.precomputed_norm_path=/norm_stats +# +# The flagship annotations are empty, so every prompt is +# "Embodiment: human bimanual. Task: fold clothes"; the HPT recipe uses the +# same default string. +defaults: + - train_zarr_cartesian + - evaluator@train_viz_evaluator: train_viz_hpt_keypoints_wrist + - override model: qwenvla_bc_mecka_kp_0p8b + - override data: mecka_fold_flagship_opsplit_hpt_keypoints + - override evaluator: eval_hpt_keypoints_wrist + - _self_ + +model: + robomimic_model: + default_prompt: "fold clothes" diff --git a/egomimic/hydra_configs/trainer/ddp_qwenvla.yaml b/egomimic/hydra_configs/trainer/ddp_qwenvla.yaml new file mode 100644 index 000000000..7c7be2414 --- /dev/null +++ b/egomimic/hydra_configs/trainer/ddp_qwenvla.yaml @@ -0,0 +1,17 @@ +# DDP for QwenVLA (Qwen 3.5 + layer-wise DiT, ~1.2 B trainable): +# find_unused_parameters: with interleaved self/cross DiT blocks the last VLM +# text layer and final norm feed no cross block, and plain DDP refuses +# parameters that get no gradient (starVLA's QwenPI has the same property). +# bf16 gradient compression: the lab's 8 x A40 nodes have no NVLink and two +# NUMA islands joined by cross-socket links, so the per-step all-reduce of +# 4.8 GB of fp32 gradients dominated the step: 6.4 s/step at 8 x 16 vs +# 1.4 s/step on one GPU (2026-09-16). Halving the bytes brought it to +# 3.6 s/step; the master weights and optimizer stay fp32. +defaults: + - ddp + +strategy: + _target_: lightning.pytorch.strategies.DDPStrategy + find_unused_parameters: true + ddp_comm_hook: + _target_: egomimic.utils.ddp_hooks.bf16_compress_hook # factory -> the torch hook diff --git a/egomimic/models/layerwise_dit.py b/egomimic/models/layerwise_dit.py new file mode 100644 index 000000000..50414479a --- /dev/null +++ b/egomimic/models/layerwise_dit.py @@ -0,0 +1,486 @@ +"""Layer-wise cross-attention DiT flow-matching head. + +Port of the GR00T N1.5 DiT as starVLA's QwenPI / QwenPI_v3 use it +(``starVLA/model/modules/action_model/flow_matching_head/cross_attention_dit.py`` +and ``LayerwiseFM_ActionHeader.py``, Apache-2.0), written with torch +primitives instead of ``diffusers``: + +* DiT block ``i`` cross-attends to VLM hidden state ``i`` (one context per + block); with ``interleave_self_attention`` the odd blocks self-attend over + the action tokens instead and their context is unused. +* Time enters every block through AdaLN (scale / shift from a sinusoidal + then MLP embedding of the CONTINUOUS flow time ``t`` in [0, 1]); a final + time-modulated norm precedes the action decoder. +* ``adaln_zero`` adds zero-initialised gates so every block is the identity + at init (user decision 3 for the deferred HPT DiT head); off for QwenVLA. + +Flow-matching convention (``egomimic/models/fm_policy.py``): +``x_t = t * noise + (1 - t) * a``, target ``noise - a``, ``t`` from 1 to 0. +""" + +from __future__ import annotations + +from typing import Optional + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from egomimic.models.denoising_nets import posemb_sincos + +# openpi's posemb_sincos period range for a flow time in [0, 1]. +TIME_MIN_PERIOD = 4e-3 +TIME_MAX_PERIOD = 4.0 + + +def _time_features(t: torch.Tensor, dim: int) -> torch.Tensor: + """(B,) flow time in [0, 1] -> (B, dim) sinusoidal features.""" + return posemb_sincos(t.float(), dim, TIME_MIN_PERIOD, TIME_MAX_PERIOD) + + +class TimestepEmbedder(nn.Module): + def __init__(self, dim: int, freq_dim: int = 256) -> None: + super().__init__() + self.freq_dim = freq_dim + self.mlp = nn.Sequential( + nn.Linear(freq_dim, dim), nn.SiLU(), nn.Linear(dim, dim) + ) + + def forward(self, t: torch.Tensor) -> torch.Tensor: + feats = _time_features(t, self.freq_dim).to(self.mlp[0].weight.dtype) + return self.mlp(feats) + + +class AdaLayerNorm(nn.Module): + """LayerNorm whose scale / shift come from the time embedding.""" + + def __init__(self, dim: int, zero_init: bool = False) -> None: + super().__init__() + self.linear = nn.Linear(dim, 2 * dim) + self.norm = nn.LayerNorm(dim, eps=1e-5, elementwise_affine=False) + if zero_init: + nn.init.zeros_(self.linear.weight) + nn.init.zeros_(self.linear.bias) + + def forward(self, x: torch.Tensor, temb: torch.Tensor) -> torch.Tensor: + scale, shift = self.linear(F.silu(temb)).chunk(2, dim=-1) + return self.norm(x) * (1 + scale[:, None]) + shift[:, None] + + +class Attention(nn.Module): + """Multi-head attention; self-attention when ``context`` is None.""" + + def __init__( + self, + dim: int, + num_heads: int, + kv_dim: Optional[int] = None, + dropout: float = 0.0, + ) -> None: + super().__init__() + if dim % num_heads: + raise ValueError(f"hidden {dim} is not divisible by {num_heads} heads") + kv_dim = dim if kv_dim is None else kv_dim + self.num_heads = num_heads + self.dropout = dropout + self.to_q = nn.Linear(dim, dim) + self.to_k = nn.Linear(kv_dim, dim) + self.to_v = nn.Linear(kv_dim, dim) + self.to_out = nn.Linear(dim, dim) + + def forward( + self, + x: torch.Tensor, + context: Optional[torch.Tensor] = None, + context_mask: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + ctx = x if context is None else context + B, T, _ = x.shape + S = ctx.shape[1] + h = self.num_heads + q = self.to_q(x).view(B, T, h, -1).transpose(1, 2) + k = self.to_k(ctx).view(B, S, h, -1).transpose(1, 2) + v = self.to_v(ctx).view(B, S, h, -1).transpose(1, 2) + attn_mask = None + if context is not None and context_mask is not None: + # (B, S) bool, True = attend -> (B, 1, 1, S), what SDPA expects + attn_mask = context_mask.to(torch.bool)[:, None, None, :] + out = F.scaled_dot_product_attention( + q, + k, + v, + attn_mask=attn_mask, + dropout_p=self.dropout if self.training else 0.0, + ) + return self.to_out(out.transpose(1, 2).reshape(B, T, -1)) + + +class DiTBlock(nn.Module): + """AdaLN -> (cross | self) attention -> LN -> feed-forward, with residuals.""" + + def __init__( + self, + dim: int, + num_heads: int, + cross_dim: Optional[int] = None, + dropout: float = 0.0, + final_dropout: bool = True, + ff_mult: int = 4, + adaln_zero: bool = False, + ) -> None: + super().__init__() + self.is_cross = cross_dim is not None + self.norm1 = AdaLayerNorm(dim, zero_init=adaln_zero) + self.attn = Attention(dim, num_heads, kv_dim=cross_dim, dropout=dropout) + self.attn_dropout = nn.Dropout(dropout) if final_dropout else nn.Identity() + self.norm2 = nn.LayerNorm(dim) + self.ff = nn.Sequential( + nn.Linear(dim, ff_mult * dim), + nn.GELU(approximate="tanh"), + nn.Dropout(dropout), + nn.Linear(ff_mult * dim, dim), + nn.Dropout(dropout) if final_dropout else nn.Identity(), + ) + self.gate = None + if adaln_zero: + self.gate = nn.Linear(dim, 2 * dim) + nn.init.zeros_(self.gate.weight) + nn.init.zeros_(self.gate.bias) + + def forward( + self, + x: torch.Tensor, + temb: torch.Tensor, + context: Optional[torch.Tensor] = None, + context_mask: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + h = self.norm1(x, temb) + if self.is_cross: + attn = self.attn(h, context, context_mask) + else: + attn = self.attn(h) + attn = self.attn_dropout(attn) + if self.gate is None: + x = x + attn + return x + self.ff(self.norm2(x)) + gate_attn, gate_ff = self.gate(F.silu(temb)).chunk(2, dim=-1) + x = x + gate_attn[:, None] * attn + return x + gate_ff[:, None] * self.ff(self.norm2(x)) + + +class LayerwiseDiT(nn.Module): + def __init__( + self, + num_layers: int, + hidden: int, + num_heads: int, + cross_dim: Optional[int] = None, + dropout: float = 0.2, + final_dropout: bool = True, + interleave_self_attention: bool = True, + adaln_zero: bool = False, + ) -> None: + super().__init__() + cross_dim = hidden if cross_dim is None else cross_dim + self.time_embed = TimestepEmbedder(hidden) + self.blocks = nn.ModuleList( + [ + DiTBlock( + hidden, + num_heads, + cross_dim=None + if (interleave_self_attention and i % 2 == 1) + else cross_dim, + dropout=dropout, + final_dropout=final_dropout, + adaln_zero=adaln_zero, + ) + for i in range(num_layers) + ] + ) + self.norm_out = AdaLayerNorm(hidden, zero_init=adaln_zero) + + def forward( + self, + tokens: torch.Tensor, + contexts: list, + context_mask: Optional[torch.Tensor], + t: torch.Tensor, + ) -> torch.Tensor: + if len(contexts) != len(self.blocks): + raise ValueError( + f"LayerwiseDiT has {len(self.blocks)} blocks but got {len(contexts)} contexts" + ) + temb = self.time_embed(t) + x = tokens + for i, block in enumerate(self.blocks): + if block.is_cross: + x = block(x, temb, contexts[i], context_mask) + else: + x = block(x, temb) + return self.norm_out(x, temb) + + +class ActionEncoder(nn.Module): + """starVLA's ActionEncoder: Linear(a) ++ sinusoid(t) -> swish(Linear) -> Linear.""" + + def __init__(self, action_dim: int, hidden: int) -> None: + super().__init__() + self.hidden = hidden + self.layer1 = nn.Linear(action_dim, hidden) + self.layer2 = nn.Linear(2 * hidden, hidden) + self.layer3 = nn.Linear(hidden, hidden) + + def forward(self, actions: torch.Tensor, t: torch.Tensor) -> torch.Tensor: + a = self.layer1(actions) + tau = ( + _time_features(t, self.hidden) + .to(a.dtype)[:, None] + .expand(-1, a.shape[1], -1) + ) + x = F.silu(self.layer2(torch.cat([a, tau], dim=-1))) + return self.layer3(x) + + +class StateEncoder(nn.Module): + """One state MLP per embodiment (starVLA's ``MLP``), one token per proprio + step, plus the HPT proprio-history extras: a learned zero-init per-step + embedding and whole-context history dropout (train only).""" + + def __init__( + self, + state_dims: dict, + hidden: int, + history_len: int = 1, + history_dropout: float = 0.0, + mlp_hidden: int = 1024, + ) -> None: + super().__init__() + self.state_dims = {str(k): int(v) for k, v in state_dims.items()} + self.history_len = int(history_len) + self.history_dropout = float(history_dropout) + self.mlps = nn.ModuleDict( + { + name: nn.Sequential( + nn.Linear(dim, mlp_hidden), nn.ReLU(), nn.Linear(mlp_hidden, hidden) + ) + for name, dim in self.state_dims.items() + } + ) + self.time_embed = ( + nn.Parameter(torch.zeros(self.history_len, hidden)) + if self.history_len > 1 + else None + ) + + def forward(self, state: torch.Tensor, embodiment: str) -> torch.Tensor: + if embodiment not in self.mlps: + raise KeyError( + f"StateEncoder has no MLP for embodiment {embodiment!r}; " + f"known: {sorted(self.mlps)}" + ) + if state.ndim != 3: + raise ValueError(f"state must be (B, K, D), got {tuple(state.shape)}") + if state.shape[1] != self.history_len: + raise ValueError( + f"state carries {state.shape[1]} history steps but history_len=" + f"{self.history_len}; set the keymap's proprio_history and the " + "model's history_len to the same K" + ) + if state.shape[-1] != self.state_dims[embodiment]: + raise ValueError( + f"{embodiment} state is {state.shape[-1]}-wide, expected " + f"{self.state_dims[embodiment]}" + ) + if self.training and self.history_dropout > 0 and self.history_len > 1: + drop = ( + torch.rand(state.shape[0], device=state.device) < self.history_dropout + ) + current = state[:, -1:].expand_as(state) + state = torch.where(drop[:, None, None], current, state) + y = self.mlps[embodiment](state) + if self.time_embed is not None: + y = y + self.time_embed[None] + return y + + +class LayerwiseFMHead(nn.Module): + """Flow-matching action head over ``LayerwiseDiT``. + + DiT sequence: ``[state x K] + [registers] + [action_horizon action tokens]``; + only the action positions are decoded. ``projectors`` map every VLM layer + to ``dit_hidden`` (identity when the widths match: QwenPI; LayerNorm + + Linear otherwise: QwenPI_v3's compressed head). + """ + + def __init__( + self, + action_width: int, + action_horizon: int, + vlm_hidden: int, + num_layers: int, + state_dims: Optional[dict] = None, + dit_hidden: Optional[int] = None, + head_dim: int = 64, + num_register_tokens: int = 32, + history_len: int = 1, + history_dropout: float = 0.0, + dropout: float = 0.2, + final_dropout: bool = True, + interleave_self_attention: bool = True, + adaln_zero: bool = False, + num_inference_steps: int = 10, + time_dist: str = "beta", + repeated_diffusion_steps: int = 1, + ) -> None: + super().__init__() + self.action_width = int(action_width) + self.action_horizon = int(action_horizon) + self.vlm_hidden = int(vlm_hidden) + self.dit_hidden = int(vlm_hidden if dit_hidden is None else dit_hidden) + self.num_inference_steps = int(num_inference_steps) + self.time_dist = time_dist + self.repeated_diffusion_steps = int(repeated_diffusion_steps) + if self.dit_hidden % head_dim: + raise ValueError( + f"dit_hidden {self.dit_hidden} not divisible by head_dim {head_dim}" + ) + if time_dist not in ("beta", "uniform"): + raise ValueError(f"time_dist must be beta or uniform, got {time_dist!r}") + + self.projectors = nn.ModuleList( + [ + nn.Identity() + if self.dit_hidden == self.vlm_hidden + else nn.Sequential( + nn.LayerNorm(self.vlm_hidden), + nn.Linear(self.vlm_hidden, self.dit_hidden), + ) + for _ in range(int(num_layers)) + ] + ) + self.action_encoder = ActionEncoder(self.action_width, self.dit_hidden) + self.pos_embed = nn.Parameter(torch.zeros(self.action_horizon, self.dit_hidden)) + nn.init.normal_(self.pos_embed, std=0.02) + self.registers = nn.Parameter( + torch.zeros(int(num_register_tokens), self.dit_hidden) + ) + nn.init.normal_(self.registers, std=0.02) + self.state_encoder = ( + StateEncoder(state_dims, self.dit_hidden, history_len, history_dropout) + if state_dims + else None + ) + self.dit = LayerwiseDiT( + int(num_layers), + self.dit_hidden, + self.dit_hidden // head_dim, + cross_dim=self.dit_hidden, + dropout=dropout, + final_dropout=final_dropout, + interleave_self_attention=interleave_self_attention, + adaln_zero=adaln_zero, + ) + self.action_decoder = nn.Sequential( + nn.Linear(self.dit_hidden, 1024), + nn.ReLU(), + nn.Linear(1024, self.action_width), + ) + + # -- pieces ------------------------------------------------------------ + + def project(self, contexts: list) -> list: + if len(contexts) != len(self.projectors): + raise ValueError( + f"got {len(contexts)} VLM layers but the head has {len(self.projectors)} projectors" + ) + return [proj(ctx) for proj, ctx in zip(self.projectors, contexts)] + + def velocity( + self, + x_t: torch.Tensor, + t: torch.Tensor, + contexts: list, + context_mask: Optional[torch.Tensor], + state: Optional[torch.Tensor] = None, + embodiment: Optional[str] = None, + ) -> torch.Tensor: + B, T, _ = x_t.shape + if T != self.action_horizon: + raise ValueError( + f"action chunk has {T} steps, head expects {self.action_horizon}" + ) + parts = [] + if state is not None: + if self.state_encoder is None: + raise ValueError( + "state given but the head was built without state_dims" + ) + parts.append(self.state_encoder(state, embodiment)) + parts.append(self.registers[None].expand(B, -1, -1)) + parts.append(self.action_encoder(x_t, t) + self.pos_embed[None]) + tokens = torch.cat(parts, dim=1) + out = self.dit(tokens, contexts, context_mask, t) + return self.action_decoder(out[:, -T:]) + + def sample_time(self, n: int, device) -> torch.Tensor: + if self.time_dist == "beta": + t = torch.distributions.Beta(1.5, 1.0).sample((n,)).to(device) + else: + t = torch.rand(n, device=device) + return t * 0.999 + 0.001 + + # -- training / inference ------------------------------------------------ + + def compute_loss( + self, + contexts: list, + context_mask: Optional[torch.Tensor], + actions: torch.Tensor, + loss_mask: torch.Tensor, + state: Optional[torch.Tensor] = None, + embodiment: Optional[str] = None, + ) -> torch.Tensor: + contexts = self.project(contexts) + R = self.repeated_diffusion_steps + if R > 1: + contexts = [c.repeat(R, 1, 1) for c in contexts] + context_mask = None if context_mask is None else context_mask.repeat(R, 1) + actions = actions.repeat(R, 1, 1) + loss_mask = loss_mask.repeat(R, 1, 1) + state = None if state is None else state.repeat(R, 1, 1) + noise = torch.randn_like(actions) + t = self.sample_time(actions.shape[0], actions.device).to(actions.dtype) + te = t[:, None, None] + x_t = te * noise + (1 - te) * actions + u_t = noise - actions + v_t = self.velocity(x_t, t, contexts, context_mask, state, embodiment) + mask = loss_mask.to(v_t.dtype).expand_as(v_t) + return ((v_t - u_t) ** 2 * mask).sum() / mask.sum().clamp(min=1.0) + + @torch.no_grad() + def sample( + self, + contexts: list, + context_mask: Optional[torch.Tensor], + state: Optional[torch.Tensor] = None, + embodiment: Optional[str] = None, + generator: Optional[torch.Generator] = None, + num_steps: Optional[int] = None, + ) -> torch.Tensor: + contexts = self.project(contexts) + B = contexts[0].shape[0] + device, dtype = contexts[0].device, self.pos_embed.dtype + n = self.num_inference_steps if num_steps is None else int(num_steps) + x = torch.randn( + (B, self.action_horizon, self.action_width), + generator=generator, + device=device, + dtype=dtype, + ) + dt = -1.0 / n + t = torch.ones(B, device=device, dtype=dtype) + for _ in range(n): + x = x + dt * self.velocity(x, t, contexts, context_mask, state, embodiment) + t = t + dt + return x diff --git a/egomimic/models/qwenvla_nets.py b/egomimic/models/qwenvla_nets.py new file mode 100644 index 000000000..d2f89f42a --- /dev/null +++ b/egomimic/models/qwenvla_nets.py @@ -0,0 +1,195 @@ +"""QwenVLA: Qwen 3.5 backbone + layer-wise cross-DiT head (starVLA QwenPI port). + +``QwenVLBackbone`` reuses ``Qwen35VLMEncoder``'s loading, freezing, dtype and +pretrained-weight bookkeeping and changes two things: the chat turn holds +every camera frame of the sample (``[image] * F + [text]``), and the forward +returns the last ``num_layers`` text-layer hidden states (one context per +DiT block) plus the padding mask, instead of the visual tokens of one layer. +State never enters the VLM (spec section 4): it goes to the head. +""" + +from __future__ import annotations + +from typing import List, Optional, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from egomimic.models.hpt_nets import Qwen35VLMEncoder +from egomimic.models.layerwise_dit import LayerwiseFMHead + + +class QwenVLBackbone(Qwen35VLMEncoder): + """Qwen 3.5 as the QwenVLA backbone. + + Args (beyond ``Qwen35VLMEncoder``): ``gradient_checkpointing`` is applied + to the VLM whenever any of it trains; ``num_layers`` (default: the text + stack's depth) is how many of the LAST hidden states are returned, i.e. + the DiT depth. ``dtype="float32"`` for a fine-tune (Lightning's bf16 + autocast does the mixed precision; AdamW must see fp32 master weights); + ``"bfloat16"`` is fine for a frozen backbone. + """ + + def __init__( + self, + model_name: str = Qwen35VLMEncoder.DEFAULT_MODEL, + dtype: str = "float32", + freeze: bool = False, + trainable_layers: int = 0, + gradient_checkpointing: bool = True, + image_size: tuple = (352, 640), + max_text_tokens: int = 128, + num_layers: Optional[int] = None, + **kwargs, + ) -> None: + super().__init__( + model_name=model_name, + dtype=dtype, + freeze=freeze, + trainable_layers=trainable_layers, + feature_layer=-1, + image_size=image_size, + max_text_tokens=max_text_tokens, + **kwargs, + ) + depth = int(self.model.config.text_config.num_hidden_layers) + self.num_layers = depth if num_layers is None else int(num_layers) + if not 1 <= self.num_layers <= depth: + raise ValueError( + f"num_layers={self.num_layers} but {self.model_name} has {depth} text layers" + ) + self.gradient_checkpointing = bool(gradient_checkpointing) and not self._no_grad + if self.gradient_checkpointing: + self.model.gradient_checkpointing_enable( + gradient_checkpointing_kwargs={"use_reentrant": False} + ) + + def backbone_parameters(self) -> List[nn.Parameter]: + """The VLM's parameters, for ``ModelWrapper._backbone_param_groups``.""" + return list(self.model.parameters()) + + def _processor_inputs(self, images: torch.Tensor, prompts: List[str]) -> dict: + """Chat layout ``[image] * F + [prompt]`` per sample; ``images`` is + ``(B, F, 3, H, W)`` or ``(B, 3, H, W)`` (F = 1).""" + if images.ndim == 4: + images = images[:, None] + batch, frames = images.shape[:2] + if len(prompts) != batch: + raise RuntimeError( + f"QwenVLBackbone got {len(prompts)} prompts for {batch} samples" + ) + tokenizer = self.processor.tokenizer + texts, flat = [], [] + for i, prompt in enumerate(prompts): + ids = tokenizer(str(prompt), add_special_tokens=False)["input_ids"] + text = tokenizer.decode(ids[: self.max_text_tokens]) + content = [{"type": "image"}] * frames + [{"type": "text", "text": text}] + texts.append( + self.processor.apply_chat_template( + [{"role": "user", "content": content}], + tokenize=False, + add_generation_prompt=False, + ) + ) + flat.extend(list(images[i])) + return self.processor( + text=texts, + images=flat, + padding=True, + return_tensors="pt", + do_rescale=False, + do_resize=False, + ) + + def _prepare_frames(self, images: torch.Tensor) -> torch.Tensor: + if images.ndim == 4: + images = images[:, None] + if images.ndim != 5 or images.shape[2] != 3: + raise ValueError( + f"images must be (B, F, 3, H, W), got {tuple(images.shape)}" + ) + batch, frames, _, height, width = images.shape + images = images.float() + if (height, width) != self.image_size: + images = F.interpolate( + images.flatten(0, 1), + size=self.image_size, + mode="bilinear", + align_corners=False, + antialias=True, + ).view(batch, frames, 3, *self.image_size) + return images.clamp(0.0, 1.0) + + def forward( + self, images: torch.Tensor, prompts: List[str] + ) -> Tuple[List[torch.Tensor], torch.Tensor]: + """One VLM forward -> (last ``num_layers`` hidden states, padding mask).""" + images = self._prepare_frames(images) + frames = images.shape[1] + inputs = self._processor_inputs(images, prompts) + device = next(self.model.parameters()).device + inputs = { + key: (value.to(device) if torch.is_tensor(value) else value) + for key, value in inputs.items() + } + if torch.is_tensor(inputs.get("pixel_values")): + inputs["pixel_values"] = inputs["pixel_values"].to(self._load_dtype) + + def run(): + return self.model.model( + **inputs, output_hidden_states=True, use_cache=False + ) + + if self._no_grad: + with torch.no_grad(): + out = run() + else: + out = run() + + counts = (inputs["input_ids"] == self.image_token_id).sum(dim=1) + expected = frames * self.num_visual_tokens + if not bool((counts == expected).all()): + raise RuntimeError( + f"expected {expected} visual tokens per sample, got {counts.tolist()}" + ) + hidden = list(out.hidden_states[-self.num_layers :]) + mask = inputs["attention_mask"].bool() + return [h.float() for h in hidden], mask + + +class QwenVLAModel(nn.Module): + """``backbone`` + ``head``; ``encoders`` is the plain dict + ``ModelWrapper._backbone_param_groups`` reads for the VLM LR group.""" + + def __init__(self, backbone: QwenVLBackbone, head: LayerwiseFMHead) -> None: + super().__init__() + self.backbone = backbone + self.head = head + self.encoders = {"vlm": backbone} + + def encode(self, data: dict): + return self.backbone(data["images"], data["prompts"]) + + def compute_loss(self, data: dict) -> torch.Tensor: + contexts, mask = self.encode(data) + return self.head.compute_loss( + contexts, + mask, + data["action"], + data["loss_mask"], + data.get("state"), + data["embodiment_name"], + ) + + def sample( + self, data: dict, generator: Optional[torch.Generator] = None + ) -> torch.Tensor: + contexts, mask = self.encode(data) + return self.head.sample( + contexts, + mask, + data.get("state"), + data["embodiment_name"], + generator=generator, + ) diff --git a/egomimic/utils/ddp_hooks.py b/egomimic/utils/ddp_hooks.py new file mode 100644 index 000000000..f6b8f9c5d --- /dev/null +++ b/egomimic/utils/ddp_hooks.py @@ -0,0 +1,14 @@ +"""DDP gradient-communication hooks, exposed as zero-argument factories so a +Hydra config can select one (``ddp_comm_hook: {_target_: ...}``). Lightning +logs the hook's ``__qualname__``, so the config must hand it the real function, +not a ``functools.partial`` (Hydra's ``_partial_``).""" + +from __future__ import annotations + +from torch.distributed.algorithms.ddp_comm_hooks import default_hooks + + +def bf16_compress_hook(): + """All-reduce gradients in bf16 (half the bytes of fp32), decompress after. + Master weights and the optimizer stay fp32.""" + return default_hooks.bf16_compress_hook diff --git a/tests/fixtures/fake_qwen35.py b/tests/fixtures/fake_qwen35.py new file mode 100644 index 000000000..a9d8da30d --- /dev/null +++ b/tests/fixtures/fake_qwen35.py @@ -0,0 +1,198 @@ +"""A fake Qwen 3.5 (multi-image chat layout) for the QwenVLA unit tests. + +Same shape of stub as ``tests/unit/test_qwen35_vlm.py`` uses for the HPT stem +pair, generalised to N images per sample: every ```` in the template +expands to one token per (2x2) patch, ``hidden_states`` are returned for +every layer, and the image-token hidden states depend on the pixels so a +wrong image/text split shows up. +""" + +from __future__ import annotations + +import json +import os +from types import SimpleNamespace + +import torch +import torch.nn as nn + +STUB_HIDDEN = 16 +STUB_LAYERS = 2 +STUB_IMAGE_SIZE = (4, 8) # (H, W) -> (4/2) * (8/2) = 8 tokens per frame +STUB_TOKENS_PER_FRAME = 8 +IMAGE_TOKEN_ID = 900 +STRIDE = 2 + + +class FakeTokenizer: + """Whitespace tokenizer with a growing vocabulary (ids start at 1).""" + + def __init__(self) -> None: + self.padding_side = "left" + self._ids: dict[str, int] = {} + self._words: dict[int, str] = {} + + def _id(self, word: str) -> int: + if word not in self._ids: + index = len(self._ids) + 1 + self._ids[word] = index + self._words[index] = word + return self._ids[word] + + def __call__(self, text, add_special_tokens=False): + return {"input_ids": [self._id(w) for w in str(text).split()]} + + def decode(self, ids): + return " ".join(self._words[i] for i in ids) + + +class FakeProcessor: + """Chat layout `` ... ``, right padded.""" + + IMAGE = "" + + def __init__(self) -> None: + self.tokenizer = FakeTokenizer() + + def apply_chat_template( + self, messages, tokenize=False, add_generation_prompt=False + ): + content = messages[0]["content"] + images = [c for c in content if c["type"] == "image"] + texts = [c["text"] for c in content if c["type"] == "text"] + assert content[: len(images)] == images, "images must come first" + return " ".join([self.IMAGE] * len(images) + texts).strip() + + def __call__( + self, + text, + images, + padding=True, + return_tensors="pt", + do_rescale=None, + do_resize=None, + ): + assert do_rescale is False and do_resize is False + images = list(images) # flat: all frames of sample 0, then sample 1, ... + n_vis = (images[0].shape[-2] // STRIDE) * (images[0].shape[-1] // STRIDE) + rows = [] + for item in text: + ids: list[int] = [] + for word in item.split(): + if word == self.IMAGE: + ids.extend([IMAGE_TOKEN_ID] * n_vis) + else: + ids.append(self.tokenizer._id(word)) + rows.append(ids) + width = max(len(r) for r in rows) + input_ids = torch.zeros(len(rows), width, dtype=torch.long) + attention = torch.zeros(len(rows), width, dtype=torch.long) + for i, row in enumerate(rows): + input_ids[i, : len(row)] = torch.tensor(row, dtype=torch.long) + attention[i, : len(row)] = 1 + return { + "input_ids": input_ids, + "attention_mask": attention, + "mm_token_type_ids": (input_ids == IMAGE_TOKEN_ID).long(), + "pixel_values": torch.stack(images).flatten(1), + } + + +class FakeInner(nn.Module): + """``model.model``: the VLM without the lm_head.""" + + def __init__(self) -> None: + super().__init__() + self.embed = nn.Embedding(1024, STUB_HIDDEN) + self.layers = nn.ModuleList( + [nn.Linear(STUB_HIDDEN, STUB_HIDDEN) for _ in range(STUB_LAYERS)] + ) + + def forward( + self, + input_ids=None, + attention_mask=None, + pixel_values=None, + output_hidden_states=False, + **kwargs, + ): + hidden = self.embed(input_ids) + if pixel_values is not None: + batch = input_ids.shape[0] + # one scalar per sample from ALL of its frames, added on its image tokens + per_sample = pixel_values.reshape(batch, -1).mean(dim=-1) + image = (input_ids == IMAGE_TOKEN_ID).unsqueeze(-1).float() + hidden = hidden + image * per_sample[:, None, None] + states = [hidden] + for layer in self.layers: + hidden = torch.tanh(layer(hidden)) + states.append(hidden) + out = {"last_hidden_state": hidden} + if output_hidden_states: + out["hidden_states"] = tuple(states) + return SimpleNamespace(**out) + + +class FakeVLM(nn.Module): + def __init__(self) -> None: + super().__init__() + self.config = SimpleNamespace( + image_token_id=IMAGE_TOKEN_ID, + text_config=SimpleNamespace( + hidden_size=STUB_HIDDEN, num_hidden_layers=STUB_LAYERS + ), + vision_config=SimpleNamespace(patch_size=STRIDE, spatial_merge_size=1), + ) + self.model = FakeInner() + self.lm_head = nn.Linear(STUB_HIDDEN, 8, bias=False) + self.lm_head.weight = nn.Parameter(torch.randn(8, STUB_HIDDEN)) + + def gradient_checkpointing_enable(self, **kwargs) -> None: + self.gradient_checkpointing = True + + +def _write_stub_snapshot(tmp_path) -> str: + from safetensors.torch import save_file + + torch.manual_seed(0) + reference = FakeVLM() + snapshot = tmp_path / "qwen35_stub" + snapshot.mkdir() + state = {k: v.contiguous() for k, v in reference.state_dict().items()} + state.pop("lm_head.weight") # tied in the real checkpoint too + state["mtp.fc.weight"] = torch.zeros(2, 2) + shard = "model.safetensors-00001-of-00001.safetensors" + save_file(state, str(snapshot / shard)) + (snapshot / "model.safetensors.index.json").write_text( + json.dumps({"weight_map": {k: shard for k in state}}), encoding="utf-8" + ) + return str(snapshot) + + +def install_fake_qwen35(tmp_path, monkeypatch) -> str: + """Write a stub snapshot and patch the two ``transformers`` entry points + ``Qwen35VLMEncoder`` (hence ``QwenVLBackbone``) loads through.""" + import transformers + from safetensors.torch import load_file + + path = _write_stub_snapshot(tmp_path) + + def _load_model(load_path, dtype=None, **kwargs): + model = FakeVLM() + state = load_file( + os.path.join(load_path, "model.safetensors-00001-of-00001.safetensors") + ) + state.pop("mtp.fc.weight", None) + model.load_state_dict(state, strict=False) + model.lm_head.weight = nn.Parameter(model.model.embed.weight[:8].clone()) + return model if dtype is None else model.to(dtype) + + monkeypatch.setattr( + transformers.AutoProcessor, "from_pretrained", lambda *a, **k: FakeProcessor() + ) + monkeypatch.setattr( + transformers.AutoModelForImageTextToText, + "from_pretrained", + staticmethod(_load_model), + ) + return path diff --git a/tests/unit/test_layerwise_dit.py b/tests/unit/test_layerwise_dit.py new file mode 100644 index 000000000..7f5e81b47 --- /dev/null +++ b/tests/unit/test_layerwise_dit.py @@ -0,0 +1,248 @@ +"""Layer-wise cross-attention DiT (GR00T N1.5 / starVLA QwenPI port).""" + +from __future__ import annotations + +import pytest +import torch +import torch.nn as nn + +from egomimic.models.layerwise_dit import DiTBlock, LayerwiseDiT + +B, N, S, H, HEADS, L = 2, 5, 7, 16, 2, 4 + + +def _dit(**kwargs) -> LayerwiseDiT: + torch.manual_seed(0) + defaults = dict(num_layers=L, hidden=H, num_heads=HEADS, dropout=0.0) + defaults.update(kwargs) + return LayerwiseDiT(**defaults).eval() + + +def _inputs(seed: int = 1): + g = torch.Generator().manual_seed(seed) + tokens = torch.randn(B, N, H, generator=g) + contexts = [torch.randn(B, S, H, generator=g) for _ in range(L)] + mask = torch.ones(B, S, dtype=torch.bool) + t = torch.rand(B, generator=g) + return tokens, contexts, mask, t + + +def test_output_shape(): + tokens, contexts, mask, t = _inputs() + assert _dit()(tokens, contexts, mask, t).shape == (B, N, H) + + +def test_interleave_alternates_cross_and_self_blocks(): + blocks = _dit(interleave_self_attention=True).blocks + assert [b.is_cross for b in blocks] == [True, False, True, False] + assert all(b.is_cross for b in _dit(interleave_self_attention=False).blocks) + + +def test_wrong_number_of_contexts_raises(): + tokens, contexts, mask, t = _inputs() + with pytest.raises(ValueError): + _dit()(tokens, contexts[:-1], mask, t) + + +def test_masked_context_positions_do_not_affect_the_output(): + dit = _dit() + tokens, contexts, mask, t = _inputs() + mask[:, -2:] = False + ref = dit(tokens, contexts, mask, t) + changed = [c.clone() for c in contexts] + for c in changed: + c[:, -2:] += 100.0 + assert torch.allclose(dit(tokens, changed, mask, t), ref, atol=1e-5) + # sanity: an unmasked change does move the output + changed[0][:, 0] += 100.0 + assert not torch.allclose(dit(tokens, changed, mask, t), ref, atol=1e-3) + + +def test_each_cross_block_reads_its_own_layer_only(): + dit = _dit(interleave_self_attention=True) + tokens, contexts, mask, t = _inputs() + ref = dit(tokens, contexts, mask, t) + # layer 1 feeds a SELF block: changing it changes nothing + c = [x.clone() for x in contexts] + c[1] += 100.0 + assert torch.allclose(dit(tokens, c, mask, t), ref, atol=1e-5) + # layer 2 feeds a cross block: changing it changes the output + c = [x.clone() for x in contexts] + c[2] += 100.0 + assert not torch.allclose(dit(tokens, c, mask, t), ref, atol=1e-3) + + +def test_adaln_zero_makes_the_stack_the_identity_up_to_the_output_norm(): + dit = _dit(adaln_zero=True) + tokens, contexts, mask, t = _inputs() + out = dit(tokens, contexts, mask, t) + expected = nn.functional.layer_norm(tokens, (H,)) + assert torch.allclose(out, expected, atol=1e-5) + + +def test_time_changes_the_output(): + dit = _dit() + tokens, contexts, mask, t = _inputs() + a = dit(tokens, contexts, mask, torch.zeros(B)) + b = dit(tokens, contexts, mask, torch.ones(B)) + assert not torch.allclose(a, b, atol=1e-3) + + +def test_block_gradients_flow_to_context(): + block = DiTBlock(H, HEADS, cross_dim=H, dropout=0.0) + x = torch.randn(B, N, H) + ctx = torch.randn(B, S, H, requires_grad=True) + temb = torch.randn(B, H) + block(x, temb, ctx, torch.ones(B, S, dtype=torch.bool)).sum().backward() + assert ctx.grad is not None and ctx.grad.abs().sum() > 0 + + +# --------------------------------------------------------------------------- +# the head +# --------------------------------------------------------------------------- +from egomimic.models.layerwise_dit import LayerwiseFMHead, StateEncoder # noqa: E402 + +W, T, VLM_H = 6, 10, 24 +STATE_DIMS = {"human_bimanual": 3, "eva_bimanual": 5} + + +def _head(**kwargs) -> LayerwiseFMHead: + torch.manual_seed(0) + defaults = dict( + action_width=W, + action_horizon=T, + vlm_hidden=VLM_H, + num_layers=L, + state_dims=STATE_DIMS, + dit_hidden=H, + head_dim=8, + num_register_tokens=3, + dropout=0.0, + num_inference_steps=3, + ) + defaults.update(kwargs) + return LayerwiseFMHead(**defaults) + + +def _head_inputs(seed: int = 2, K: int = 1): + g = torch.Generator().manual_seed(seed) + contexts = [torch.randn(B, S, VLM_H, generator=g) for _ in range(L)] + mask = torch.ones(B, S, dtype=torch.bool) + actions = torch.randn(B, T, W, generator=g) + loss_mask = torch.ones(B, T, W) + state = torch.randn(B, K, STATE_DIMS["human_bimanual"], generator=g) + return contexts, mask, actions, loss_mask, state + + +def test_head_loss_is_finite_and_backpropagates(): + head = _head().train() + contexts, mask, actions, loss_mask, state = _head_inputs() + loss = head.compute_loss( + contexts, mask, actions, loss_mask, state, "human_bimanual" + ) + assert torch.isfinite(loss) + loss.backward() + assert head.projectors[0][1].weight.grad is not None + assert head.state_encoder.mlps["human_bimanual"][0].weight.grad is not None + + +def test_projectors_are_identity_when_widths_match(): + same = _head(vlm_hidden=H, dit_hidden=H) + assert all(isinstance(p, nn.Identity) for p in same.projectors) + assert all(isinstance(p, nn.Sequential) for p in _head().projectors) + + +def test_loss_is_the_masked_mean_of_the_velocity_error(): + """The mask weights the LOSS only (the noised input still carries every + dim), so check the reduction against a seeded replica of compute_loss.""" + head = _head().eval() + contexts, mask, actions, loss_mask, state = _head_inputs() + loss_mask[..., -2:] = 0.0 + torch.manual_seed(5) + loss = head.compute_loss( + contexts, mask, actions, loss_mask, state, "human_bimanual" + ) + torch.manual_seed(5) # same draw order: noise first, then t + noise = torch.randn_like(actions) + t = head.sample_time(B, actions.device).to(actions.dtype) + te = t[:, None, None] + v = head.velocity( + te * noise + (1 - te) * actions, + t, + head.project(contexts), + mask, + state, + "human_bimanual", + ) + expected = ((v - (noise - actions)) ** 2)[..., :-2].mean() + assert torch.allclose(loss, expected, atol=1e-6) + + +def test_sample_shape_and_determinism_under_a_generator(): + head = _head().eval() + contexts, mask, _, _, state = _head_inputs() + g1 = torch.Generator().manual_seed(11) + g2 = torch.Generator().manual_seed(11) + a = head.sample(contexts, mask, state, "human_bimanual", generator=g1) + b = head.sample(contexts, mask, state, "human_bimanual", generator=g2) + assert a.shape == (B, T, W) + assert torch.equal(a, b) + g3 = torch.Generator().manual_seed(12) + assert not torch.equal( + a, head.sample(contexts, mask, state, "human_bimanual", generator=g3) + ) + + +def test_state_selects_the_embodiment_mlp_and_checks_its_width(): + head = _head().eval() + contexts, mask, _, _, state = _head_inputs() + g = torch.Generator().manual_seed(1) + with_state = head.sample(contexts, mask, state, "human_bimanual", generator=g) + g = torch.Generator().manual_seed(1) + without = head.sample(contexts, mask, None, "human_bimanual", generator=g) + assert not torch.allclose(with_state, without) + with pytest.raises(ValueError): + head.sample(contexts, mask, state, "eva_bimanual") # 3-wide state, 5-wide MLP + with pytest.raises(KeyError): + head.sample(contexts, mask, state, "aria_bimanual") + + +def test_history_time_embed_is_zero_init_and_history_length_is_checked(): + head = _head(history_len=2) + assert head.state_encoder.time_embed.shape == (2, H) + assert torch.all(head.state_encoder.time_embed == 0) + assert _head(history_len=1).state_encoder.time_embed is None + contexts, mask, actions, loss_mask, state = _head_inputs(K=1) + with pytest.raises(ValueError): + head.compute_loss(contexts, mask, actions, loss_mask, state, "human_bimanual") + + +def test_history_dropout_replaces_the_past_with_the_current_step(): + enc = StateEncoder(STATE_DIMS, H, history_len=3, history_dropout=1.0).train() + state = torch.randn(B, 3, 3) + current_only = state[:, -1:].expand_as(state) + assert torch.allclose( + enc(state, "human_bimanual"), enc(current_only, "human_bimanual") + ) + enc.eval() + assert not torch.allclose( + enc(state, "human_bimanual"), enc(current_only, "human_bimanual") + ) + + +def test_repeated_diffusion_steps_repeats_the_batch(): + head = _head(repeated_diffusion_steps=2).train() + contexts, mask, actions, loss_mask, state = _head_inputs() + loss = head.compute_loss( + contexts, mask, actions, loss_mask, state, "human_bimanual" + ) + assert torch.isfinite(loss) + + +def test_action_horizon_mismatch_raises(): + head = _head() + contexts, mask, actions, loss_mask, state = _head_inputs() + with pytest.raises(ValueError): + head.compute_loss( + contexts, mask, actions[:, :-1], loss_mask[:, :-1], state, "human_bimanual" + ) diff --git a/tests/unit/test_pi05_norm_rot6d.py b/tests/unit/test_pi05_norm_rot6d.py index e90b91115..0e07d75c4 100644 --- a/tests/unit/test_pi05_norm_rot6d.py +++ b/tests/unit/test_pi05_norm_rot6d.py @@ -502,9 +502,10 @@ def test_fallback_cap_returns_last_good_sample_or_raises(caplog): assert any("last good sample" in r.message for r in caplog.records) -def test_video_fps_compensates_for_world_size(): - # Distributed val strides an episode's frames by world_size on each rank; - # playback fps must scale down to keep videos wall-clock real-time. +def test_video_fps_is_the_source_rate_regardless_of_world_size(): + """Val loaders are CombinedLoaders that Lightning never shards, so rank 0 + renders every frame: playback must stay at the source rate (an earlier + world_size division made 8-GPU videos 8x slow motion).""" from egomimic.eval.eval_video import EvalVideo class _Stub(EvalVideo): @@ -512,11 +513,12 @@ def compute_metrics_and_viz(self, batch, do_viz=True): raise NotImplementedError ev = _Stub.__new__(_Stub) - for world, expected in [(1, 30), (2, 15), (4, 8), (8, 4)]: + for world in (1, 2, 8): ev.trainer = SimpleNamespace(world_size=world) - assert ev._video_fps() == expected, (world, ev._video_fps()) - ev.trainer = SimpleNamespace() # no world_size attr -> assume 1 + assert ev._video_fps() == 30, (world, ev._video_fps()) + ev.trainer = SimpleNamespace(world_size=None) assert ev._video_fps() == 30 + assert ev._video_fps(source_fps=10) == 10 def test_viz_gate_follows_lightning_epoch_convention(): diff --git a/tests/unit/test_qwenvla_algo.py b/tests/unit/test_qwenvla_algo.py new file mode 100644 index 000000000..3a166c83c --- /dev/null +++ b/tests/unit/test_qwenvla_algo.py @@ -0,0 +1,209 @@ +"""QwenVLA algo on the fake VLM: prompts, batch conversion, train / eval paths.""" + +from __future__ import annotations + +import functools + +import pytest +import torch +from fixtures.fake_qwen35 import STUB_IMAGE_SIZE, install_fake_qwen35 + +from egomimic.algo.qwenvla import QwenVLA +from egomimic.models.layerwise_dit import LayerwiseFMHead +from egomimic.models.qwenvla_nets import QwenVLBackbone +from egomimic.rldb.embodiment.embodiment import get_embodiment_id + +EMB = "human_bimanual" +EMB_ID = get_embodiment_id(EMB) +CAM = "observations.images.front_img_1" +PROPRIO = "observations.state.ee_pose" +AC = "actions_cartesian" +B, T, D_ACT, D_PROP, WIDTH = 2, 4, 6, 3, 8 + + +class _StubNormStats: + """The four norm-stat calls the algo makes, for one embodiment.""" + + def keys_of_type(self, kind, embodiment_id): + return { + "camera_keys": [CAM], + "proprio_keys": [PROPRIO], + "action_keys": [AC], + "lang_keys": [], + }[kind] + + def is_key_with_embodiment(self, key, embodiment_id): + return key in (CAM, PROPRIO, AC) + + def zarr_key_to_keyname(self, key, embodiment_id): + return key + + def unnormalize(self, predictions, embodiment_id): + return {k: v * 2.0 for k, v in predictions.items()} + + +def _algo(snapshot, history_len: int = 1, **kwargs) -> QwenVLA: + torch.manual_seed(0) + backbone = QwenVLBackbone( + model_name=snapshot, dtype="float32", image_size=STUB_IMAGE_SIZE + ) + head = functools.partial( + LayerwiseFMHead, + action_horizon=T, + head_dim=8, + num_register_tokens=2, + dropout=0.0, + num_inference_steps=2, + ) + defaults = dict( + norm_stats=_StubNormStats(), + backbone=backbone, + head=head, + domains=[EMB], + dims={EMB: {"proprio": D_PROP, "action": D_ACT}}, + ac_keys={EMB: AC}, + action_width=WIDTH, + annotation_key="annotations", + default_prompt="", + history_len=history_len, + device=torch.device("cpu"), + ) + defaults.update(kwargs) + return QwenVLA(**defaults) + + +def _raw_batch(K: int = 1, annotations=None) -> dict: + g = torch.Generator().manual_seed(1) + proprio = torch.randn(B, K, D_PROP, generator=g) + if K == 1: + proprio = proprio[:, 0] + batch = { + AC: torch.randn(B, T, D_ACT, generator=g), + PROPRIO: proprio, + CAM: torch.rand(B, 3, *STUB_IMAGE_SIZE, generator=g), + "action_pad_mask": torch.tensor( + [[1, 1, 1, 0], [1, 1, 1, 1]], dtype=torch.float32 + ), + "annotations": annotations + if annotations is not None + else [["fold the towel"], []], + } + return {EMB: batch} + + +@pytest.fixture() +def snapshot(tmp_path, monkeypatch): + return install_fake_qwen35(tmp_path, monkeypatch) + + +def test_compose_prompt(snapshot): + algo = _algo(snapshot) + assert ( + algo.compose_prompt(EMB, "fold the towel") + == "Embodiment: human bimanual. Task: fold the towel" + ) + assert algo.compose_prompt(EMB, "") == "Embodiment: human bimanual." + plain = _algo(snapshot, embodiment_label=False) + assert plain.compose_prompt(EMB, "fold") == "Task: fold" + modes = _algo(snapshot, control_mode={"human": "wrist frame"}) + assert ( + modes.compose_prompt(EMB, "") + == "Embodiment: human bimanual. Control mode: wrist frame." + ) + + +def test_process_batch_builds_prompts_pad_mask_and_embodiment(snapshot): + algo = _algo(snapshot) + algo.nets.train() + out = algo.process_batch_for_training(_raw_batch())[EMB_ID] + assert out["sampled_prompt"] == ["fold the towel", ""] + assert out["pad_mask"].shape == (B, T, 1) and torch.all(out["pad_mask"] == 1) + assert out["embodiment"].tolist() == [EMB_ID] + masked = _algo(snapshot, use_pad_mask=True).process_batch_for_training( + _raw_batch() + )[EMB_ID] + assert masked["pad_mask"][0, -1, 0] == 0 and masked["pad_mask"][1, -1, 0] == 1 + + +def test_model_data_pads_the_action_and_masks_the_padding(snapshot): + algo = _algo(snapshot) + batch = algo.process_batch_for_training(_raw_batch())[EMB_ID] + data = algo._to_model_data(batch, EMB_ID) + assert data["action"].shape == (B, T, WIDTH) + assert torch.all(data["action"][..., D_ACT:] == 0) + assert data["loss_mask"].shape == (B, T, WIDTH) + assert torch.all(data["loss_mask"][..., :D_ACT] == 1) + assert torch.all(data["loss_mask"][..., D_ACT:] == 0) + assert data["images"].shape == (B, 1, 3, *STUB_IMAGE_SIZE) + assert data["state"].shape == (B, 1, D_PROP) + assert data["prompts"] == [ + "Embodiment: human bimanual. Task: fold the towel", + "Embodiment: human bimanual.", + ] + assert data["embodiment_name"] == EMB + + +def test_forward_training_returns_a_finite_loss(snapshot): + algo = _algo(snapshot) + algo.nets.train() + batch = algo.process_batch_for_training(_raw_batch()) + preds = algo.forward_training(batch) + assert torch.isfinite(preds[f"{EMB}_loss"]) + losses = algo.compute_losses(preds, batch) + assert losses["action_loss"] == preds[f"{EMB}_loss"] + log = algo.log_info({"losses": losses}) + assert log["Loss"] == losses["action_loss"].item() + + +def test_forward_eval_is_deterministic_and_native_width(snapshot): + algo = _algo(snapshot) + algo.nets.eval() + + def one_pass(): + algo.reset_eval_pass_counter() + batch = algo.process_batch_for_training(_raw_batch()) + return algo.forward_eval(batch) + + a, b = one_pass(), one_pass() + assert a[f"{EMB}_{AC}"].shape == (B, T, D_ACT) + assert torch.equal(a[f"{EMB}_{AC}"], b[f"{EMB}_{AC}"]) + assert torch.equal(a[f"{EMB}_loss"], b[f"{EMB}_loss"]) + # a second batch within one pass draws different noise + algo.reset_eval_pass_counter() + first = algo.forward_eval(algo.process_batch_for_training(_raw_batch())) + second = algo.forward_eval(algo.process_batch_for_training(_raw_batch())) + assert not torch.equal(first[f"{EMB}_{AC}"], second[f"{EMB}_{AC}"]) + + +def test_proprio_history_feeds_k_state_tokens(snapshot): + algo = _algo(snapshot, history_len=2) + algo.nets.train() + batch = algo.process_batch_for_training(_raw_batch(K=2)) + assert torch.isfinite(algo.forward_training(batch)[f"{EMB}_loss"]) + with pytest.raises(ValueError): + algo.forward_training(algo.process_batch_for_training(_raw_batch(K=1))) + + +def test_action_wider_than_action_width_is_rejected(snapshot): + with pytest.raises(ValueError): + _algo(snapshot, action_width=D_ACT - 1) + + +def test_backbone_parameters_go_to_the_vlm_lr_group(snapshot): + algo = _algo(snapshot) + vlm = algo.nets["policy"].encoders["vlm"] + assert vlm.backbone_parameters() and all( + p.requires_grad for p in vlm.backbone_parameters() + ) + + +def test_algo_is_built_on_the_cpu_and_the_device_setter_moves_the_nets(snapshot): + """Every DDP rank constructs the algo before Lightning assigns its GPU; the + nets must not be parked on cuda:0 meanwhile (8 x 4.8 GB overflowed it on a + checkpoint resume, 2026-09-16).""" + algo = _algo(snapshot) + assert algo.device == torch.device("cpu") + assert next(algo.nets.parameters()).device.type == "cpu" + algo.device = "cpu" # what ModelWrapper.on_fit_start writes (a str or device) + assert algo.device == torch.device("cpu") + assert next(algo.nets.parameters()).device.type == "cpu" diff --git a/tests/unit/test_qwenvla_configs.py b/tests/unit/test_qwenvla_configs.py new file mode 100644 index 000000000..750362cf7 --- /dev/null +++ b/tests/unit/test_qwenvla_configs.py @@ -0,0 +1,208 @@ +"""QwenVLA configs compose, wire the algo, and train two stub steps.""" + +from __future__ import annotations + +import pytest +from fixtures.fake_qwen35 import STUB_HIDDEN, STUB_IMAGE_SIZE, install_fake_qwen35 +from hydra import compose, initialize_config_module + +MODEL = "qwenvla_bc_mecka_6d_0p8b" + + +def _compose(config_name: str, overrides: list[str]): + with initialize_config_module( + config_module="egomimic.hydra_configs", version_base=None + ): + return compose(config_name=config_name, overrides=overrides) + + +def test_0p8b_config_wires_the_algo(): + rm = _compose("train_zarr_cartesian", [f"model={MODEL}"]).model.robomimic_model + assert rm._target_ == "egomimic.algo.qwenvla.QwenVLA" + assert rm.backbone._target_ == "egomimic.models.qwenvla_nets.QwenVLBackbone" + assert rm.backbone.model_name == "Qwen/Qwen3.5-0.8B" + assert rm.backbone.freeze is False and rm.backbone.dtype == "float32" + assert rm.backbone.image_size == [352, 640] + assert rm.head._target_ == "egomimic.models.layerwise_dit.LayerwiseFMHead" + assert rm.head._partial_ is True + assert rm.head.action_horizon == 100 + assert rm.head.dit_hidden is None + assert rm.head.interleave_self_attention is True + assert rm.head.repeated_diffusion_steps == 1 + # the HPT baseline head samples with 50 steps; the metric compares sampled actions + assert rm.head.num_inference_steps == 50 + assert rm.dims.human_bimanual == {"proprio": 20, "action": 18} + assert rm.action_width == 18 + assert rm.embodiment_label is True + assert rm.history_len == 1 + + +def test_2b_config_compresses_the_head(): + rm = _compose( + "train_zarr_cartesian", ["model=qwenvla_bc_mecka_6d_2b"] + ).model.robomimic_model + assert rm.backbone.model_name == "Qwen/Qwen3.5-2B" + assert rm.head.dit_hidden == 1024 + assert rm.backbone.freeze is True and rm.backbone.trainable_layers == 8 + + +def test_cotrain_config_shares_one_padded_action_width(): + rm = _compose( + "train_zarr_cartesian", ["model=qwenvla_cotrain_eva_mecka_6d"] + ).model.robomimic_model + assert list(rm.domains) == ["eva_bimanual", "human_bimanual"] + assert rm.action_width == 20 + assert rm.dims.eva_bimanual.action == 20 and rm.dims.human_bimanual.action == 18 + + +def test_flagship_recipe_selects_qwenvla_and_the_hpt_evaluator(): + cfg = _compose("train_zarr_mecka_flagship_6d_qwenvla", []) + assert cfg.model.robomimic_model._target_ == "egomimic.algo.qwenvla.QwenVLA" + assert cfg.evaluator._target_ == "egomimic.eval.eval_hpt.HPTEvalVideo" + assert cfg.model.backbone_lr_scale == 0.1 + assert cfg.model.optimizer.lr == 1e-4 + + +KP_MODEL = "qwenvla_bc_mecka_kp_0p8b" + + +def test_kp_config_is_the_0p8b_on_the_144d_keypoint_action(): + rm = _compose("train_zarr_cartesian", [f"model={KP_MODEL}"]).model.robomimic_model + assert rm.dims.human_bimanual == {"proprio": 144, "action": 144} + assert rm.ac_keys.human_bimanual == "actions_keypoints" + assert rm.action_width == 144 + assert rm["6dof"] is False + # everything else is inherited from the 6D config + assert rm.backbone.model_name == "Qwen/Qwen3.5-0.8B" + assert rm.head.num_inference_steps == 50 + + +@pytest.mark.parametrize( + "recipe, target", + [ + ("train_zarr_mecka_flagship_kp_qwenvla", "egomimic.algo.qwenvla.QwenVLA"), + ("train_zarr_mecka_flagship_kp_hpt", "egomimic.algo.hpt.HPT"), + ], +) +def test_kp_recipes_pair_the_keypoint_data_evaluator_and_prompt(recipe, target): + cfg = _compose(recipe, []) + rm = cfg.model.robomimic_model + assert rm._target_ == target + assert rm.ac_keys.human_bimanual == "actions_keypoints" + assert rm.dims.human_bimanual.action == 144 + # both arms of the comparison see the same prompt (annotations are empty) + assert rm.default_prompt == "fold clothes" + assert cfg.evaluator._target_ == "egomimic.eval.eval_hpt.HPTEvalVideo" + assert ( + "keypoints_revert_6d_wristframe" + in cfg.evaluator.transform_lists.human_bimanual._target_ + ) + resolver = cfg.data.train_datasets.human_bimanual.resolver + assert resolver.key_map.keymap_mode == "keypoints" + assert resolver.transform_list.mode == "keypoints_wristframe_6d" + assert resolver.transform_list.pad_proprio_gripper is False + + +def test_kp_top3_data_keeps_the_three_operators_and_the_keypoint_pipeline(): + cfg = _compose( + "train_zarr_mecka_flagship_kp_hpt", + ["data=mecka_fold_flagship_top3_hpt_keypoints"], + ) + lambdas = cfg.data.train_datasets.human_bimanual.filters.filter_lambdas + assert any("6903686e0e94ce070afd1f24" in fn for fn in lambdas) + assert ( + cfg.data.train_datasets.human_bimanual.resolver.transform_list.mode + == "keypoints_wristframe_6d" + ) + + +def test_hpt_kp_model_carries_the_history_and_backbone_lr_knobs(): + cfg = _compose("train_zarr_cartesian", ["model=hpt_bc_keypoints_wrist_300M"]) + stem = cfg.model.robomimic_model.stem_specs.human_bimanual.state_keypoints + assert stem.history_len == 1 and stem.history_dropout == 0.2 + assert cfg.model.backbone_lr_scale == 1.0 + + +def test_qwenvla_train_step_on_the_stub(tmp_path, monkeypatch): + from fixtures.recipes import Recipe, common_overrides, cpu_trainer_overrides + from fixtures.train_harness import compose_recipe, hermetic_env, write_fixtures + + import egomimic.trainHydra as train_hydra + + snapshot = install_fake_qwen35(tmp_path, monkeypatch) + hermetic_env(monkeypatch) + recipe = Recipe("train_zarr_cartesian", "mecka_all_6d", MODEL, "human_bimanual") + data, out, hashes = write_fixtures(tmp_path, "mecka") + cfg = compose_recipe( + recipe, + common_overrides( + recipe.embodiment, + data, + out, + batch_size=2, + num_workers=0, + episode_hashes=hashes, + ) + + cpu_trainer_overrides(2) + + [ + f"model.robomimic_model.backbone.model_name={snapshot}", + "model.robomimic_model.backbone.dtype=float32", + f"model.robomimic_model.backbone.image_size=[{STUB_IMAGE_SIZE[0]},{STUB_IMAGE_SIZE[1]}]", + "model.robomimic_model.head.head_dim=8", + "model.robomimic_model.head.num_register_tokens=2", + "model.robomimic_model.head.num_inference_steps=2", + ], + out, + ) + metrics, objects = train_hydra.train(cfg) + loss = float(metrics["Train/action_loss"]) + assert loss == loss, "non-finite loss" + assert objects["trainer"].global_step == 2 + # objects["model"] is the ModelWrapper (trainHydra.py object_dict) + policy = objects["model"].model.nets["policy"] + assert policy.head.dit_hidden == STUB_HIDDEN + assert len(policy.head.projectors) == policy.backbone.num_layers + + +def test_qwenvla_kp_train_step_on_the_stub(tmp_path, monkeypatch): + from fixtures.recipes import Recipe, common_overrides, cpu_trainer_overrides + from fixtures.train_harness import compose_recipe, hermetic_env, write_fixtures + + import egomimic.trainHydra as train_hydra + + snapshot = install_fake_qwen35(tmp_path, monkeypatch) + hermetic_env(monkeypatch) + recipe = Recipe("train_zarr_cartesian", "mecka", KP_MODEL, "human_bimanual") + data, out, hashes = write_fixtures(tmp_path, "mecka") + emb = "data.train_datasets.human_bimanual.resolver" + cfg = compose_recipe( + recipe, + common_overrides( + recipe.embodiment, + data, + out, + batch_size=2, + num_workers=0, + episode_hashes=hashes, + ) + + cpu_trainer_overrides(2) + + [ + f"{emb}.key_map.include_ee_pose=false", + f"{emb}.transform_list.include_ee_pose=false", + f"{emb}.transform_list.pad_proprio_gripper=false", + f"model.robomimic_model.backbone.model_name={snapshot}", + "model.robomimic_model.backbone.dtype=float32", + f"model.robomimic_model.backbone.image_size=[{STUB_IMAGE_SIZE[0]},{STUB_IMAGE_SIZE[1]}]", + "model.robomimic_model.head.head_dim=8", + "model.robomimic_model.head.num_register_tokens=2", + "model.robomimic_model.head.num_inference_steps=2", + ], + out, + ) + metrics, objects = train_hydra.train(cfg) + loss = float(metrics["Train/action_loss"]) + assert loss == loss, "non-finite loss" + assert objects["trainer"].global_step == 2 + policy = objects["model"].model.nets["policy"] + assert policy.head.action_width == 144 + assert policy.head.state_encoder.state_dims == {"human_bimanual": 144} diff --git a/tests/unit/test_qwenvla_nets.py b/tests/unit/test_qwenvla_nets.py new file mode 100644 index 000000000..04ea1a9e2 --- /dev/null +++ b/tests/unit/test_qwenvla_nets.py @@ -0,0 +1,141 @@ +"""QwenVLBackbone (Qwen 3.5, multi-frame, all layers) + QwenVLAModel on a fake VLM.""" + +from __future__ import annotations + +import pytest +import torch +from fixtures.fake_qwen35 import ( + STUB_HIDDEN, + STUB_IMAGE_SIZE, + STUB_LAYERS, + STUB_TOKENS_PER_FRAME, + install_fake_qwen35, +) + +from egomimic.models.hpt_nets import verify_pretrained_weights +from egomimic.models.layerwise_dit import LayerwiseFMHead +from egomimic.models.qwenvla_nets import QwenVLAModel, QwenVLBackbone + +B, T, W = 2, 4, 6 + + +@pytest.fixture() +def snapshot(tmp_path, monkeypatch): + return install_fake_qwen35(tmp_path, monkeypatch) + + +def _backbone(snapshot, **kwargs) -> QwenVLBackbone: + kwargs.setdefault("image_size", STUB_IMAGE_SIZE) + kwargs.setdefault("dtype", "float32") + return QwenVLBackbone(model_name=snapshot, **kwargs) + + +def _frames(frames: int = 1, seed: int = 0) -> torch.Tensor: + g = torch.Generator().manual_seed(seed) + return torch.rand(B, frames, 3, *STUB_IMAGE_SIZE, generator=g) + + +def test_backbone_returns_every_text_layer_and_the_padding_mask(snapshot): + backbone = _backbone(snapshot).eval() + prompts = ["Embodiment: human bimanual. Task: fold", "Embodiment: human bimanual."] + contexts, mask = backbone(_frames(), prompts) + assert len(contexts) == STUB_LAYERS == backbone.num_layers + assert all(c.shape == (B, mask.shape[1], STUB_HIDDEN) for c in contexts) + assert mask.dtype == torch.bool + # 8 image tokens + 5 words vs 8 + 3 words: right padding shows in the mask + assert mask[0].sum() == STUB_TOKENS_PER_FRAME + 5 + assert mask[1].sum() == STUB_TOKENS_PER_FRAME + 3 + assert mask[1, -2:].tolist() == [False, False] + + +def test_backbone_accepts_a_single_frame_batch_and_multi_frame(snapshot): + backbone = _backbone(snapshot).eval() + one, mask1 = backbone(_frames()[:, 0], ["a b"] * B) + two, mask2 = backbone(_frames(2), ["a b"] * B) + assert mask1.shape[1] == STUB_TOKENS_PER_FRAME + 2 + assert mask2.shape[1] == 2 * STUB_TOKENS_PER_FRAME + 2 + + +def test_backbone_resizes_frames_to_image_size(snapshot): + backbone = _backbone(snapshot).eval() + big = torch.rand(B, 1, 3, 8, 16) + contexts, mask = backbone(big, ["x"] * B) + assert mask.shape[1] == STUB_TOKENS_PER_FRAME + 1 + + +def test_num_layers_takes_the_last_layers(snapshot): + backbone = _backbone(snapshot, num_layers=1).eval() + full = _backbone(snapshot).eval() + contexts, _ = backbone(_frames(), ["x"] * B) + all_layers, _ = full(_frames(), ["x"] * B) + assert len(contexts) == 1 + assert torch.allclose(contexts[0], all_layers[-1]) + with pytest.raises(ValueError): + _backbone(snapshot, num_layers=STUB_LAYERS + 1) + + +def test_image_tokens_depend_on_the_pixels(snapshot): + backbone = _backbone(snapshot).eval() + a, _ = backbone(_frames(seed=0), ["x"] * B) + b, _ = backbone(_frames(seed=1), ["x"] * B) + assert not torch.allclose( + a[-1][:, :STUB_TOKENS_PER_FRAME], b[-1][:, :STUB_TOKENS_PER_FRAME] + ) + + +def test_fine_tune_exposes_trainable_backbone_parameters(snapshot): + tuned = _backbone(snapshot, freeze=False) + params = tuned.backbone_parameters() + assert params and all(p.requires_grad for p in params) + assert getattr(tuned.model, "gradient_checkpointing", False) is True + frozen = _backbone(snapshot, freeze=True) + assert not any(p.requires_grad for p in frozen.backbone_parameters()) + contexts, _ = frozen.eval()(_frames(), ["x"] * B) + assert not contexts[-1].requires_grad + + +def test_backbone_passes_the_weight_hash_check(snapshot): + verify_pretrained_weights(_backbone(snapshot), verbose=False) + + +def _model(snapshot, **head_kwargs) -> QwenVLAModel: + backbone = _backbone(snapshot) + torch.manual_seed(0) + head = LayerwiseFMHead( + action_width=W, + action_horizon=T, + vlm_hidden=backbone.hidden_size, + num_layers=backbone.num_layers, + state_dims={"human_bimanual": 3}, + head_dim=8, + num_register_tokens=2, + dropout=0.0, + num_inference_steps=2, + **head_kwargs, + ) + return QwenVLAModel(backbone, head) + + +def _data() -> dict: + g = torch.Generator().manual_seed(3) + return { + "images": _frames(), + "prompts": ["Embodiment: human bimanual. Task: fold"] * B, + "action": torch.randn(B, T, W, generator=g), + "loss_mask": torch.ones(B, T, W), + "state": torch.randn(B, 1, 3, generator=g), + "embodiment_name": "human_bimanual", + } + + +def test_model_loss_and_sample(snapshot): + model = _model(snapshot) + assert model.encoders["vlm"] is model.backbone + model.train() + loss = model.compute_loss(_data()) + assert torch.isfinite(loss) + loss.backward() + model.eval() + g = torch.Generator().manual_seed(0) + actions = model.sample(_data(), generator=g) + assert actions.shape == (B, T, W)