File size: 4,458 Bytes
1e05592 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 | #!/usr/bin/env python3
"""
Pre-compute and cache belief vectors for all policy label windows.
Since the SFT backbone is fully frozen, belief[i] = SFTModel(frames[i]) is
deterministic. Computing it once and saving it eliminates the 3B-param VLM
forward pass from every training step, making PolicyHead training ~1000× faster.
Output:
data/belief_cache/train.pt — tensors for all train samples
data/belief_cache/val.pt — tensors for all val samples
Cache format (per split):
{
"beliefs": FloatTensor [N, hidden_dim] (float32)
"tta_means": FloatTensor [N]
"tta_vars": FloatTensor [N]
}
Indices match exactly the sample order in data/policy_labels/{split}.json.
Usage:
cd PROJECT_ROOT
python -m training.Policy.make_belief_cache \
--sft_checkpoint checkpoints/SFT/sft_v2/best \
--label_dir data/policy_labels \
--out_dir data/belief_cache \
--batch_size 8
"""
from __future__ import annotations
import argparse
import json
import logging
from pathlib import Path
from typing import List
import torch
from torch.utils.data import DataLoader
from tqdm import tqdm
import sys
sys.path.insert(0, str(Path(__file__).resolve().parent.parent.parent))
from .policy_model import PolicyModel
from .policy_dataset import PolicyDataset, policy_collate_fn
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
logger = logging.getLogger("Policy.make_cache")
@torch.no_grad()
def build_cache(
model: PolicyModel,
loader: DataLoader,
split_name: str,
) -> dict:
"""Run VLM on all samples, collect belief + tta statistics."""
model.eval()
all_beliefs: List[torch.Tensor] = []
all_tta_means: List[torch.Tensor] = []
all_tta_vars: List[torch.Tensor] = []
for batch in tqdm(loader, desc=f" Caching {split_name}"):
inputs = model._build_inputs(batch["images"], batch["metadata"])
from torch.amp import autocast
with autocast(device_type="cuda", dtype=model._amp_dtype, enabled=True):
belief = model.sft.encode_observation(inputs)
tta_mean, tta_logvar = model.sft.tta_head(belief)
tta_var = torch.exp(tta_logvar.float().clamp(-20.0, 20.0))
all_beliefs.append(belief.float().cpu())
all_tta_means.append(tta_mean.float().cpu())
all_tta_vars.append(tta_var.cpu())
beliefs = torch.cat(all_beliefs, dim=0)
tta_means = torch.cat(all_tta_means, dim=0)
tta_vars = torch.cat(all_tta_vars, dim=0)
logger.info(
f" {split_name}: cached {beliefs.shape[0]} samples "
f"belief shape={tuple(beliefs.shape)} "
f"size={beliefs.nbytes / 1e6:.1f} MB"
)
return {"beliefs": beliefs, "tta_means": tta_means, "tta_vars": tta_vars}
def main():
parser = argparse.ArgumentParser("make_belief_cache")
parser.add_argument("--sft_checkpoint", required=True)
parser.add_argument("--label_dir", default="data/policy_labels")
parser.add_argument("--out_dir", default="data/belief_cache")
parser.add_argument("--batch_size", type=int, default=8,
help="Larger = faster caching (no grad, more GPU memory)")
parser.add_argument("--splits", nargs="+", default=["train", "val"])
args = parser.parse_args()
odir = Path(args.out_dir)
odir.mkdir(parents=True, exist_ok=True)
logger.info("Loading SFTModel (frozen backbone for belief extraction)...")
model = PolicyModel(args.sft_checkpoint, use_bf16=True)
for split in args.splits:
label_path = Path(args.label_dir) / f"{split}.json"
if not label_path.exists():
logger.warning(f" {label_path} not found — skipping {split}")
continue
out_path = odir / f"{split}.pt"
if out_path.exists():
logger.info(f" Cache already exists: {out_path} — skipping")
continue
logger.info(f"\nBuilding cache for split: {split}")
ds = PolicyDataset([label_path], split=split)
loader = DataLoader(
ds, batch_size=args.batch_size, shuffle=False,
num_workers=4, collate_fn=policy_collate_fn, pin_memory=True,
)
cache = build_cache(model, loader, split)
torch.save(cache, out_path)
logger.info(f" Saved → {out_path}")
logger.info("\n✅ Belief cache complete.")
if __name__ == "__main__":
main()
|