| """Dense reward: one 1-3 action-quality score per 0.5 s, with rationales for bad spans. |
| |
| python -m batch_utils.dense_reward --jobs jobs.jsonl --output dense.jsonl |
| python -m batch_utils.dense_reward --video top.mp4 wrist.mp4 --task "Put the cup on the tray." --output dense.jsonl |
| |
| Each job is ``{"id": ..., "task": ..., "videos": [camera videos, in collage order]}``. |
| Scores: 1 = failure, 2 = suboptimal, 3 = normal. The episode is scored in overlapping 16 s |
| clips that are stitched into one score per source frame. Scoring is retrospective (a clip |
| sees 4 s past each kept moment), so use it for offline labeling, not as an online signal. |
| """ |
|
|
| from __future__ import annotations |
|
|
| from collections import Counter |
|
|
| import numpy as np |
|
|
| from . import core |
| from .engine import ( |
| COLLAGE_VIDEO_SIZE, |
| base_parser, |
| load_llm, |
| load_processor, |
| log, |
| open_output, |
| read_jobs, |
| render, |
| resolve_model, |
| write_record, |
| ) |
|
|
|
|
| def summarize(job, collage, windows, answers, texts, tail_trim_s, include_frames): |
| scores, clip_of_frame, group_of_frame = core.stitch_dense(windows, answers, collage.num_frames, tail_trim_s) |
| counts = Counter(scores.tolist()) |
| record = { |
| "id": job["id"], |
| "task": job["task"], |
| "fps": collage.fps, |
| "num_frames": collage.num_frames, |
| "duration_s": round(collage.duration_s, 3), |
| "worst_score": int(scores.min()), |
| "mean_score": round(float(scores.mean()), 4), |
| "fraction": {core.LABELS[s]: round(counts.get(s, 0) / collage.num_frames, 4) for s in (1, 2, 3)}, |
| "segments": core.dense_segments(scores, clip_of_frame, group_of_frame, answers, collage.fps), |
| |
| "per_second": [int(scores[min(collage.num_frames - 1, int((t + 0.5) * collage.fps))]) |
| for t in range(int(np.ceil(collage.duration_s)))], |
| "clips": [ |
| { |
| "start_s": round(window.start_s, 3), |
| "end_s": round(window.end_s, 3), |
| "quality": "".join(map(str, answer.scores)), |
| "regions": core.window_regions(window, answer), |
| "valid": answer.valid, |
| "output": text, |
| } |
| for window, answer, text in zip(windows, answers, texts) |
| ], |
| "invalid_clips": sum(not a.valid for a in answers), |
| } |
| if include_frames: |
| record["frame_scores"] = scores.tolist() |
| return record |
|
|
|
|
| def main(argv=None) -> int: |
| p = base_parser("FLORA rewarder: dense action-quality scores, 1-3 per 0.5 s.") |
| p.add_argument("--tail-trim", type=float, default=core.DENSE_TAIL_TRIM_S, |
| help="Seconds dropped from the end of every clip but the last when stitching.") |
| p.add_argument("--frame-scores", action="store_true", help="Also write one score per source frame.") |
| p.add_argument("--max-model-len", type=int, default=16384) |
| args = p.parse_args(argv) |
| jobs = read_jobs(args) |
| model_path = resolve_model(args.model, args.revision, weights=not args.dry_run) |
| processor = load_processor(model_path) |
|
|
| if args.dry_run: |
| job = jobs[0] |
| collage = core.Collage(job["videos"]) |
| windows = core.dense_windows(collage.num_frames, collage.fps) |
| prompt = render(processor, core.dense_messages(job["task"], windows[0]), thinking=False) |
| log(f"{job['id']}: {collage.num_frames} frames at {collage.fps:g} FPS, collage " |
| f"{collage.width}x{collage.height}, {len(windows)} clips\n{prompt}") |
| return 0 |
|
|
| from vllm import SamplingParams |
| from vllm.sampling_params import StructuredOutputsParams |
|
|
| llm = load_llm(model_path, args, max_model_len=args.max_model_len, |
| mm_processor_kwargs={"do_sample_frames": False, "size": COLLAGE_VIDEO_SIZE}) |
| params = SamplingParams(temperature=0.0, max_tokens=core.DENSE_MAX_TOKENS, |
| structured_outputs=StructuredOutputsParams(regex=core.dense_regex())) |
| with open_output(args) as stream: |
| for number, job in enumerate(jobs, 1): |
| collage = core.Collage(job["videos"]) |
| windows = core.dense_windows(collage.num_frames, collage.fps) |
| collage.frames(sorted({i for window in windows for i in window.frame_indices})) |
| texts = [] |
| for start in range(0, len(windows), args.max_num_seqs): |
| batch = windows[start : start + args.max_num_seqs] |
| requests = [ |
| { |
| "prompt": render(processor, core.dense_messages(job["task"], window), thinking=False), |
| "multi_modal_data": { |
| "video": ( |
| collage.frames(window.frame_indices), |
| core.video_metadata(window.frame_indices, collage.fps, collage.num_frames), |
| ) |
| }, |
| "mm_processor_kwargs": {"do_sample_frames": False}, |
| } |
| for window in batch |
| ] |
| texts += [out.outputs[0].text for out in llm.generate(requests, params, use_tqdm=False)] |
| answers = [core.parse_dense(text) for text in texts] |
| write_record(stream, summarize(job, collage, windows, answers, texts, args.tail_trim, args.frame_scores)) |
| collage.release() |
| log(f"dense: {number}/{len(jobs)} episodes ({len(windows)} clips)") |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
