scripts/smoke_phase21a_robometer_low.py

#!/usr/bin/env python3
"""One isolated Robometer-Low smoke run on real Route-A preference pairs."""

from __future__ import annotations

import argparse
import json
import math
import os
import random
import sys
import time
from functools import lru_cache
from pathlib import Path

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from PIL import Image


EXPERIMENT_ROOT = Path(__file__).resolve().parents[1]
if str(EXPERIMENT_ROOT) not in sys.path:
    sys.path.insert(0, str(EXPERIMENT_ROOT))

from src.data_contract import letterbox_rgb  # noqa: E402


SPECIAL_TOKENS = (
    "<|split_token|>",
    "<|reward_token|>",
    "<|pref_token|>",
    "<|sim_token|>",
    "<|prog_token|>",
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--snapshot", type=Path, required=True)
    parser.add_argument("--manifest", type=Path, required=True)
    parser.add_argument("--processed-root", type=Path, required=True)
    parser.add_argument("--micro-batch-pairs", type=int, required=True)
    parser.add_argument("--gradient-accumulation", type=int, required=True)
    parser.add_argument("--optimizer-steps", type=int, default=1)
    parser.add_argument("--pair-pool-size", type=int, default=8)
    parser.add_argument("--seed", type=int, default=7)
    parser.add_argument("--head-lr", type=float, default=1.0e-4)
    parser.add_argument("--lora-lr", type=float, default=2.0e-6)
    parser.add_argument("--output", type=Path, required=True)
    return parser.parse_args()


class RobometerHeads(nn.Module):
    """Exact Robometer prediction-head topology, with 10-bin progress."""

    def __init__(self, hidden_size: int, dropout: float = 0.1) -> None:
        super().__init__()
        self.progress_head = self._head(hidden_size, 10, dropout)
        self.success_head = self._head(hidden_size, 1, dropout)
        self.preference_head = self._head(hidden_size, 1, dropout)

    @staticmethod
    def _head(hidden_size: int, output_size: int, dropout: float) -> nn.Sequential:
        return nn.Sequential(
            nn.Linear(hidden_size, hidden_size // 2),
            nn.LayerNorm(hidden_size // 2),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_size // 2, output_size),
        )


def read_train_rows(path: Path, count: int) -> list[dict]:
    candidates = []
    with path.open("r", encoding="utf-8") as stream:
        for line in stream:
            row = json.loads(line)
            if row["split"] == "train" and row["sample_type"] == "preference":
                candidates.append(row)
    # Capacity must cover the worst text-padding case in the formal manifest.
    rows = sorted(candidates, key=lambda row: len(row["prompt_task"]), reverse=True)[:count]
    if len(rows) < count:
        raise RuntimeError(f"requested {count} smoke rows, found {len(rows)}")
    if not any(row["chosen"]["progress_supervised"] for row in rows):
        raise RuntimeError("smoke pool has no progress-supervised chosen trajectory")
    if not any(any(row["chosen"]["success_label"]) for row in rows):
        raise RuntimeError("smoke pool has no positive success label")
    return rows


def resolve_frame_file(value: str, processed_root: Path) -> Path:
    path = Path(value)
    if path.is_absolute():
        return path
    resolved = (processed_root.parent / path).resolve()
    if not resolved.is_file():
        raise FileNotFoundError(resolved)
    return resolved


@lru_cache(maxsize=64)
def load_trajectory_cached(
    frame_path_string: str, selected_indices: tuple[int, ...]
) -> tuple[Image.Image, ...]:
    with np.load(frame_path_string) as archive:
        frames = archive["frames"][np.asarray(selected_indices, dtype=np.int64)]
    return tuple(Image.fromarray(letterbox_rgb(frame, 256)) for frame in frames)


def load_images(record: dict, processed_root: Path) -> tuple[Image.Image, ...]:
    frame_path = resolve_frame_file(record["source_frame_file"], processed_root)
    return load_trajectory_cached(
        str(frame_path), tuple(int(index) for index in record["selected_frame_indices"])
    )


def build_conversation(
    row: dict, processed_root: Path, chosen_is_a: bool
) -> tuple[list[dict], dict, float]:
    record_a = row["chosen"] if chosen_is_a else row["rejected"]
    record_b = row["rejected"] if chosen_is_a else row["chosen"]
    prompt = (
        f"Given these two trajectories for the task '{row['prompt_task']}', evaluate "
        "which one makes more progress towards the task. Return A for the first "
        "trajectory and B for the second trajectory. Also predict the task progress "
        "at each frame of the first trajectory, how far along the robot is towards "
        "completing the task, a float between 0 and 1, where 0 is the starting state "
        "and 1 is when the task is completed. If the robot is not performing the "
        "same task, predict 0 progress."
    )
    content: list[dict] = [
        {"type": "text", "text": prompt},
        {"type": "text", "text": "This is Trajectory A. "},
    ]
    for image in load_images(record_a, processed_root):
        content.append({"type": "image", "image": image})
        content.append({"type": "text", "text": "<|prog_token|>"})
    content.extend(
        [
            {"type": "text", "text": "<|split_token|>"},
            {"type": "text", "text": "This is Trajectory B. "},
        ]
    )
    for image in load_images(record_b, processed_root):
        content.append({"type": "image", "image": image})
        content.append({"type": "text", "text": "<|prog_token|>"})
    content.extend(
        [
            {
                "type": "text",
                "text": "Now predict the preference between the two trajectories.",
            },
            {"type": "text", "text": "<|pref_token|>"},
        ]
    )
    return [{"role": "user", "content": content}], record_a, float(chosen_is_a)


def process_batch(processor, conversations: list[list[dict]]) -> dict[str, torch.Tensor]:
    from qwen_vl_utils import process_vision_info

    texts = [
        processor.apply_chat_template(
            message,
            tokenize=False,
            add_generation_prompt=False,
            add_vision_id=True,
            enable_thinking=False,
            fps=1,
        )
        for message in conversations
    ]
    image_inputs, video_inputs, video_kwargs = process_vision_info(
        conversations,
        return_video_kwargs=True,
        return_video_metadata=True,
        image_patch_size=processor.image_processor.patch_size,
    )
    if video_inputs:
        raise RuntimeError("Robometer-Low smoke requires multi-image mode, not video mode")
    kwargs = dict(
        text=texts,
        images=image_inputs,
        padding=True,
        truncation=False,
        return_tensors="pt",
        do_resize=False,
    )
    if video_kwargs:
        kwargs.update(video_kwargs)
    return processor(**kwargs)


def make_targets(records_a: list[dict], device: torch.device) -> dict[str, torch.Tensor]:
    progress = torch.as_tensor(
        np.asarray([record["target_progress"] for record in records_a]),
        dtype=torch.float32,
        device=device,
    )
    progress_mask = torch.as_tensor(
        np.asarray(
            [
                np.asarray(record["frame_valid_mask"], dtype=bool)
                & (np.asarray(record["predict_last_frame_mask"]) > 0)
                & bool(record["progress_supervised"])
                for record in records_a
            ]
        ),
        dtype=torch.bool,
        device=device,
    )
    success = torch.as_tensor(
        np.asarray([record["success_label"] for record in records_a]),
        dtype=torch.float32,
        device=device,
    )
    success_mask = torch.as_tensor(
        np.asarray([record["frame_valid_mask"] for record in records_a]),
        dtype=torch.bool,
        device=device,
    )
    return {
        "progress": progress,
        "progress_mask": progress_mask,
        "success": success,
        "success_mask": success_mask,
    }


def finite_grad_audit(named_parameters) -> tuple[bool, float, int]:
    squared_norm = 0.0
    tensor_count = 0
    for _name, parameter in named_parameters:
        if not parameter.requires_grad or parameter.grad is None:
            continue
        tensor_count += 1
        gradient = parameter.grad.detach().float()
        if not torch.isfinite(gradient).all():
            return False, math.nan, tensor_count
        squared_norm += float(torch.sum(gradient * gradient))
    return True, math.sqrt(squared_norm), tensor_count


def group_grad_norm(named_parameters, marker: str) -> float:
    squared_norm = 0.0
    found = 0
    for name, parameter in named_parameters:
        if marker not in name or parameter.grad is None:
            continue
        squared_norm += float(torch.sum(parameter.grad.detach().float() ** 2))
        found += 1
    return math.sqrt(squared_norm) if found else 0.0


def main() -> None:
    args = parse_args()
    if args.micro_batch_pairs <= 0 or args.gradient_accumulation <= 0:
        raise ValueError("batch and accumulation must be positive")
    os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed_all(args.seed)
    torch.backends.cuda.matmul.allow_tf32 = True

    from peft import LoraConfig, get_peft_model
    from transformers import AutoModelForImageTextToText, AutoProcessor

    started = time.perf_counter()
    processor = AutoProcessor.from_pretrained(
        args.snapshot, local_files_only=True, padding_side="right"
    )
    original_vocab = len(processor.tokenizer)
    for token in SPECIAL_TOKENS:
        if token not in processor.tokenizer.get_vocab():
            processor.tokenizer.add_special_tokens(
                {"additional_special_tokens": [token]}
            )
    special_token_ids = {
        token: int(processor.tokenizer.convert_tokens_to_ids(token))
        for token in SPECIAL_TOKENS
    }
    model = AutoModelForImageTextToText.from_pretrained(
        args.snapshot,
        local_files_only=True,
        dtype=torch.bfloat16,
        attn_implementation="flash_attention_2",
        device_map={"": 0},
    )
    model.resize_token_embeddings(len(processor.tokenizer))
    for parameter in model.parameters():
        parameter.requires_grad_(False)

    language_model = model.model.language_model
    peft_config = LoraConfig(
        r=8,
        lora_alpha=16,
        lora_dropout=0.05,
        bias="none",
        target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
        trainable_token_indices=list(special_token_ids.values()),
    )
    model.model.language_model = get_peft_model(language_model, peft_config)
    model.config.use_cache = False
    model.gradient_checkpointing_enable(
        gradient_checkpointing_kwargs={"use_reentrant": False}
    )
    model.enable_input_require_grads()

    hidden_size = int(model.config.text_config.hidden_size)
    heads = RobometerHeads(hidden_size).to(device="cuda", dtype=torch.bfloat16)
    backbone_trainable = [
        parameter for parameter in model.parameters() if parameter.requires_grad
    ]
    head_parameters = list(heads.parameters())
    optimizer = torch.optim.AdamW(
        [
            {"params": backbone_trainable, "lr": args.lora_lr},
            {"params": head_parameters, "lr": args.head_lr},
        ],
        weight_decay=0.01,
    )
    rows = read_train_rows(args.manifest, args.pair_pool_size)
    total_parameters = sum(parameter.numel() for parameter in model.parameters())
    trainable_backbone_parameters = sum(p.numel() for p in backbone_trainable)
    trainable_head_parameters = sum(p.numel() for p in head_parameters)

    model.train()
    heads.train()
    torch.cuda.empty_cache()
    torch.cuda.reset_peak_memory_stats()
    load_seconds = time.perf_counter() - started
    run_started = time.perf_counter()
    component_history = []
    sequence_lengths = []
    optimizer.zero_grad(set_to_none=True)
    named_trainable = list(model.named_parameters()) + [
        (f"heads.{name}", parameter) for name, parameter in heads.named_parameters()
    ]

    for optimizer_step in range(args.optimizer_steps):
        for accumulation_index in range(args.gradient_accumulation):
            global_microstep = (
                optimizer_step * args.gradient_accumulation + accumulation_index
            )
            batch_rows = [
                rows[(global_microstep * args.micro_batch_pairs + offset) % len(rows)]
                for offset in range(args.micro_batch_pairs)
            ]
            conversations = []
            records_a = []
            preference_labels = []
            for offset, row in enumerate(batch_rows):
                chosen_is_a = ((global_microstep * args.micro_batch_pairs + offset) % 2) == 0
                conversation, record_a, label = build_conversation(
                    row, args.processed_root, chosen_is_a
                )
                conversations.append(conversation)
                records_a.append(record_a)
                preference_labels.append(label)
            inputs = process_batch(processor, conversations)
            sequence_lengths.extend(inputs["attention_mask"].sum(1).tolist())
            inputs = {key: value.to("cuda") for key, value in inputs.items()}
            targets = make_targets(records_a, torch.device("cuda"))
            preference_target = torch.tensor(
                preference_labels, dtype=torch.float32, device="cuda"
            )
            output = model.model(**inputs, output_hidden_states=False, use_cache=False)
            hidden = output.last_hidden_state
            prog_mask = inputs["input_ids"].eq(special_token_ids["<|prog_token|>"])
            pref_mask = inputs["input_ids"].eq(special_token_ids["<|pref_token|>"])
            prog_counts = prog_mask.sum(1)
            pref_counts = pref_mask.sum(1)
            if not torch.equal(
                prog_counts, torch.full_like(prog_counts, 16)
            ) or not torch.equal(pref_counts, torch.ones_like(pref_counts)):
                raise RuntimeError(
                    f"token extraction mismatch: prog={prog_counts.tolist()} "
                    f"pref={pref_counts.tolist()}"
                )
            prog_hidden = hidden[prog_mask].reshape(args.micro_batch_pairs, 16, hidden_size)
            pref_hidden = hidden[pref_mask].reshape(args.micro_batch_pairs, hidden_size)
            progress_logits = heads.progress_head(prog_hidden[:, :8])
            success_logits = heads.success_head(prog_hidden[:, :8]).squeeze(-1)
            preference_logits = heads.preference_head(pref_hidden).squeeze(-1)

            if targets["progress_mask"].any():
                per_frame_progress = -(
                    targets["progress"] * F.log_softmax(progress_logits.float(), dim=-1)
                ).sum(-1)
                progress_loss = per_frame_progress[targets["progress_mask"]].mean()
            else:
                progress_loss = progress_logits.float().sum() * 0.0
            valid_success = targets["success"][targets["success_mask"]]
            valid_success_logits = success_logits.float()[targets["success_mask"]]
            positives = valid_success.sum()
            negatives = valid_success.numel() - positives
            pos_weight = (negatives / positives.clamp_min(1.0)).clamp_min(1.0)
            success_loss = F.binary_cross_entropy_with_logits(
                valid_success_logits, valid_success, pos_weight=pos_weight
            )
            preference_loss = F.binary_cross_entropy_with_logits(
                preference_logits.float(), preference_target
            )
            total_loss = progress_loss + success_loss + preference_loss
            if not torch.isfinite(total_loss):
                raise FloatingPointError(f"non-finite loss: {float(total_loss)}")
            (total_loss / args.gradient_accumulation).backward()
            component_history.append(
                {
                    "total": float(total_loss.detach()),
                    "preference": float(preference_loss.detach()),
                    "progress": float(progress_loss.detach()),
                    "success": float(success_loss.detach()),
                    "progress_supervised_frames": int(targets["progress_mask"].sum()),
                    "success_supervised_frames": int(targets["success_mask"].sum()),
                }
            )

        finite_gradients, preclip_grad_norm, gradient_tensor_count = finite_grad_audit(
            named_trainable
        )
        if not finite_gradients:
            raise FloatingPointError("non-finite trainable gradient")
        grad_groups = {
            "lora": group_grad_norm(named_trainable, "lora_"),
            "token_rows": group_grad_norm(named_trainable, "trainable_tokens"),
            "progress_head": group_grad_norm(named_trainable, "heads.progress_head"),
            "success_head": group_grad_norm(named_trainable, "heads.success_head"),
            "preference_head": group_grad_norm(named_trainable, "heads.preference_head"),
        }
        missing_grad_groups = [name for name, value in grad_groups.items() if not value > 0]
        if missing_grad_groups:
            raise RuntimeError(f"missing gradient groups: {missing_grad_groups}")
        torch.nn.utils.clip_grad_norm_(backbone_trainable + head_parameters, 1.0)
        optimizer.step()
        optimizer.zero_grad(set_to_none=True)

    torch.cuda.synchronize()
    run_seconds = time.perf_counter() - run_started
    device = torch.cuda.get_device_properties(0)
    peak_allocated = torch.cuda.max_memory_allocated()
    peak_reserved = torch.cuda.max_memory_reserved()
    effective_batch = args.micro_batch_pairs * args.gradient_accumulation
    payload = {
        "status": "pass",
        "protocol_id": "phase21a_robometer_low_smoke_v1",
        "scope": {
            "smoke_only": True,
            "formal_training_started": False,
            "cosmos3_edge_modified": False,
        },
        "micro_batch_pairs": args.micro_batch_pairs,
        "gradient_accumulation": args.gradient_accumulation,
        "effective_batch_pairs": effective_batch,
        "optimizer_steps": args.optimizer_steps,
        "microsteps": args.optimizer_steps * args.gradient_accumulation,
        "real_pair_pool_size": len(rows),
        "pair_pool_selection": "longest_prompt_characters_in_train_split",
        "pair_pool_prompt_character_range": [
            min(len(row["prompt_task"]) for row in rows),
            max(len(row["prompt_task"]) for row in rows),
        ],
        "frames_per_pair": 16,
        "precision": "BF16",
        "gradient_checkpointing": True,
        "architecture": {
            "joint_a_b_forward": True,
            "direct_preference_token_head": True,
            "per_frame_progress_tokens": True,
            "progress_bins": 10,
            "lora_rank": 8,
            "lora_alpha": 16,
            "lora_targets": ["q_proj", "k_proj", "v_proj", "o_proj"],
            "special_tokens": special_token_ids,
            "original_vocab_size": original_vocab,
            "resized_vocab_size": len(processor.tokenizer),
        },
        "parameters": {
            "base_total": total_parameters,
            "backbone_trainable_lora_and_token_rows": trainable_backbone_parameters,
            "heads_trainable": trainable_head_parameters,
            "all_trainable": trainable_backbone_parameters + trainable_head_parameters,
        },
        "loss_last": component_history[-1],
        "loss_all_finite": all(
            math.isfinite(value)
            for item in component_history
            for key, value in item.items()
            if key in {"total", "preference", "progress", "success"}
        ),
        "gradient_audit": {
            "finite": finite_gradients,
            "preclip_global_norm": preclip_grad_norm,
            "tensor_count": gradient_tensor_count,
            "group_norms": grad_groups,
        },
        "sequence_tokens": {
            "min": int(min(sequence_lengths)),
            "max": int(max(sequence_lengths)),
            "mean": float(np.mean(sequence_lengths)),
        },
        "timing": {
            "model_load_seconds": load_seconds,
            "run_seconds": run_seconds,
            "seconds_per_microstep": run_seconds
            / (args.optimizer_steps * args.gradient_accumulation),
            "pairs_per_second": effective_batch * args.optimizer_steps / run_seconds,
        },
        "gpu": {
            "name": device.name,
            "total_gib": device.total_memory / 2**30,
            "peak_allocated_gib": peak_allocated / 2**30,
            "peak_reserved_gib": peak_reserved / 2**30,
            "peak_reserved_fraction": peak_reserved / device.total_memory,
        },
    }
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8")
    print(json.dumps(payload, indent=2), flush=True)


if __name__ == "__main__":
    main()