lerobot/scripts/overfit_vf_variant.py

"""Overfit any distributional VF architecture on a small real-data batch."""

from __future__ import annotations

import argparse

import numpy as np
import torch

from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
from lerobot.datasets import LeRobotDataset, LeRobotDatasetMetadata
from lerobot.datasets.factory import resolve_delta_timestamps
from lerobot.rewards.distributional_value_function.configuration_distributional_value_function import (
    DistributionalVFConfig,
)
from lerobot.rewards.distributional_value_function.processor_distributional_value_function import (
    IMAGE_MASK_SUFFIX,
)
from lerobot.rewards.factory import make_reward_model, make_reward_pre_post_processors
from lerobot.rewards.nanovlm_value_function.configuration_nanovlm_value_function import (
    NanoVLMVFConfig,
)
from lerobot.rewards.nanovlm_value_function.processor_nanovlm_value_function import (
    NANOVLM_ATTENTION_MASK,
    NANOVLM_IMAGES,
    NANOVLM_INPUT_IDS,
)
from lerobot.rewards.temporal_siglip_value_function.configuration_temporal_siglip_value_function import (
    TemporalSiglipVFConfig,
)
from lerobot.utils.constants import OBS_LANGUAGE_ATTENTION_MASK, OBS_LANGUAGE_TOKENS, OBS_STATE


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--dataset_repo_id", required=True)
    parser.add_argument("--root", default=None)
    parser.add_argument(
        "--reward_type",
        choices=(
            "distributional_value_function",
            "temporal_siglip_value_function",
            "nanovlm_value_function",
        ),
        required=True,
    )
    parser.add_argument("--vlm_pretrained_path", default=None)
    parser.add_argument("--nanovlm_pretrained_path", default="lusxvr/nanoVLM-460M-8k")
    parser.add_argument("--num_samples", type=int, default=16)
    parser.add_argument("--steps", type=int, default=500)
    parser.add_argument("--lr_head", type=float, default=1e-3)
    parser.add_argument("--lr_backbone", type=float, default=1e-5)
    parser.add_argument("--history_steps", type=int, default=6)
    parser.add_argument("--history_frame_gap", type=int, default=30)
    parser.add_argument("--log_every", type=int, default=25)
    args = parser.parse_args()

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    metadata = LeRobotDatasetMetadata(args.dataset_repo_id, root=args.root)
    input_features = {
        key: PolicyFeature(type=FeatureType.VISUAL, shape=tuple(metadata.features[key]["shape"]))
        for key in metadata.camera_keys
    }
    if OBS_STATE in metadata.features:
        input_features[OBS_STATE] = PolicyFeature(
            type=FeatureType.STATE,
            shape=tuple(metadata.features[OBS_STATE]["shape"]),
        )
    common = {"input_features": input_features, "device": str(device), "target_method": "dirac_delta"}
    if args.reward_type == "distributional_value_function":
        config = DistributionalVFConfig(
            **common,
            vlm_pretrained_path=args.vlm_pretrained_path,
            freeze_vision_encoder=True,
        )
    elif args.reward_type == "temporal_siglip_value_function":
        config = TemporalSiglipVFConfig(
            **common,
            history_steps=args.history_steps,
            frame_gap=args.history_frame_gap,
        )
        config.normalization_mapping = {
            "VISUAL": NormalizationMode.IDENTITY,
            "STATE": NormalizationMode.MEAN_STD,
        }
    else:
        config = NanoVLMVFConfig(
            **common,
            nanovlm_pretrained_path=args.nanovlm_pretrained_path,
        )

    delta_timestamps = resolve_delta_timestamps(config, metadata)
    dataset = LeRobotDataset(
        args.dataset_repo_id,
        root=args.root,
        delta_timestamps=delta_timestamps,
        video_backend="pyav",
    )
    indices = np.linspace(0, len(dataset) - 1, args.num_samples, dtype=int).tolist()
    preprocessor, _ = make_reward_pre_post_processors(
        config,
        dataset_stats=metadata.stats,
    )

    processed_samples = []
    returns = []
    terminals = []
    for index in indices:
        sample = dataset[index]
        returns.append(torch.as_tensor(sample["mc_return"]).reshape(-1)[0])
        terminals.append(torch.as_tensor(sample["is_terminal"]).reshape(-1)[0])
        processed_samples.append(preprocessor(sample))
    batch = _collate_processed(processed_samples)
    batch["mc_return"] = torch.stack(returns).to(device)
    batch["is_terminal"] = torch.stack(terminals).bool().to(device)

    model = make_reward_model(config).to(device)
    _run_stage(
        model,
        batch,
        steps=args.steps // 2,
        learning_rate=args.lr_head,
        head_only=True,
        log_every=args.log_every,
        label="head probe",
    )
    _run_stage(
        model,
        batch,
        steps=args.steps - args.steps // 2,
        learning_rate=args.lr_backbone,
        head_only=False,
        log_every=args.log_every,
        label="fine-tune",
    )
    _image_shuffle_diagnostic(model, batch, metadata.camera_keys)


def _collate_processed(samples):
    if NANOVLM_IMAGES in samples[0]:
        max_length = max(sample[NANOVLM_INPUT_IDS].shape[1] for sample in samples)
        input_ids = []
        attention_masks = []
        for sample in samples:
            padding = max_length - sample[NANOVLM_INPUT_IDS].shape[1]
            input_ids.append(torch.nn.functional.pad(sample[NANOVLM_INPUT_IDS], (padding, 0)))
            attention_masks.append(torch.nn.functional.pad(sample[NANOVLM_ATTENTION_MASK], (padding, 0)))
        return {
            NANOVLM_IMAGES: [sample[NANOVLM_IMAGES][0] for sample in samples],
            NANOVLM_INPUT_IDS: torch.cat(input_ids),
            NANOVLM_ATTENTION_MASK: torch.cat(attention_masks),
        }

    keys = {
        *[key for key in samples[0] if key.startswith("observation.images.")],
        OBS_LANGUAGE_TOKENS,
        OBS_LANGUAGE_ATTENTION_MASK,
    }
    if OBS_STATE in samples[0]:
        keys.add(OBS_STATE)
    for key in list(keys):
        if key.startswith("observation.images.") and not key.endswith(IMAGE_MASK_SUFFIX):
            keys.add(key + IMAGE_MASK_SUFFIX)
    return {key: torch.cat([sample[key] for sample in samples], dim=0) for key in keys}


def _set_trainable(model, *, head_only: bool):
    for param in model.parameters():
        param.requires_grad = False
    for param in model.value_head.parameters():
        param.requires_grad = True
    if hasattr(model, "value_query"):
        for param in model.value_query.parameters():
            param.requires_grad = True
    if head_only:
        return

    if hasattr(model, "multi_modal_projector"):
        model.multi_modal_projector.requires_grad_(True)
        model.language_model.requires_grad_(True)
    elif hasattr(model, "temporal_transformer"):
        for name, param in model.named_parameters():
            if not name.startswith("siglip."):
                param.requires_grad = True
    else:
        model.nanovlm.MP.requires_grad_(True)
        model.nanovlm.decoder.requires_grad_(True)


def _run_stage(model, batch, *, steps, learning_rate, head_only, log_every, label):
    _set_trainable(model, head_only=head_only)
    model.train()
    params = [param for param in model.parameters() if param.requires_grad]
    optimizer = torch.optim.AdamW(params, lr=learning_rate)
    print(f"\n{label}: {sum(param.numel() for param in params):,} trainable parameters")
    for step in range(steps + 1):
        optimizer.zero_grad(set_to_none=True)
        loss, metrics = model(batch)
        loss.backward()
        optimizer.step()
        if step % log_every == 0 or step == steps:
            print(
                f"step={step:04d} loss={metrics['loss']:.4f} "
                f"mae={metrics['mae']:.4f} acc={metrics['acc_neighbor']:.3f}"
            )


@torch.no_grad()
def _image_shuffle_diagnostic(model, batch, camera_keys):
    model.eval()
    matched_loss, _ = model(batch)
    shuffled = dict(batch)
    permutation = torch.roll(torch.arange(batch["mc_return"].shape[0], device=matched_loss.device), 1)
    if NANOVLM_IMAGES in batch:
        shuffled[NANOVLM_IMAGES] = [batch[NANOVLM_IMAGES][index] for index in permutation.cpu().tolist()]
    for key in camera_keys:
        if key in batch:
            shuffled[key] = batch[key][permutation]
    shuffled_loss, _ = model(shuffled)
    print(
        f"\nvisual dependence: matched_loss={matched_loss.item():.4f} "
        f"shuffled_loss={shuffled_loss.item():.4f} "
        f"gap={shuffled_loss.item() - matched_loss.item():+.4f}"
    )


if __name__ == "__main__":
    main()