#!/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()