| """Terminal reward: one 1-5 success score for each whole episode. |
| |
| python -m batch_utils.terminal_reward --jobs jobs.jsonl --output terminal.jsonl |
| python -m batch_utils.terminal_reward --video top.mp4 wrist.mp4 --task "Put the cup on the tray." --output terminal.jsonl |
| |
| Each job is ``{"id": ..., "task": ..., "videos": [camera videos, in collage order]}``. |
| """ |
|
|
| from __future__ import annotations |
|
|
| 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 main(argv=None) -> int: |
| p = base_parser("FLORA rewarder: terminal (episode success) score, 1-5.") |
| p.add_argument("--with-rationale", action="store_true", help="Ask for a one-line rationale before the score.") |
| 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) |
|
|
| prepared = [] |
| for job in jobs: |
| collage = core.Collage(job["videos"]) |
| indices = core.terminal_frame_indices(collage.num_frames, collage.fps) |
| messages = core.terminal_messages(job["task"], args.with_rationale) |
| prepared.append((job, collage, indices, render(processor, messages, thinking=False))) |
| if args.dry_run: |
| job, collage, indices, prompt = prepared[0] |
| log(f"{job['id']}: {collage.num_frames} frames at {collage.fps:g} FPS, " |
| f"collage {collage.width}x{collage.height}, {len(indices)} sampled frames\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}, |
| max_logprobs=core.TERMINAL_TOP_LOGPROBS) |
| tokenizer = llm.get_tokenizer() |
| params = SamplingParams( |
| temperature=0.0, |
| max_tokens=core.terminal_max_tokens(args.with_rationale), |
| logprobs=core.TERMINAL_TOP_LOGPROBS, |
| structured_outputs=StructuredOutputsParams(regex=core.terminal_regex(args.with_rationale)), |
| ) |
| with open_output(args) as stream: |
| for start in range(0, len(prepared), args.max_num_seqs): |
| batch = prepared[start : start + args.max_num_seqs] |
| requests = [ |
| { |
| "prompt": prompt, |
| "multi_modal_data": { |
| "video": (collage.frames(indices), core.video_metadata(indices, collage.fps, collage.num_frames)) |
| }, |
| "mm_processor_kwargs": {"do_sample_frames": False}, |
| } |
| for _, collage, indices, prompt in batch |
| ] |
| outputs = llm.generate(requests, params, use_tqdm=False) |
| for _, collage, _, _ in batch: |
| collage.release() |
| for (job, collage, indices, _), output in zip(batch, outputs): |
| completion = output.outputs[0] |
| score, rationale = core.parse_terminal(completion.text, args.with_rationale) |
| positions = [ |
| ( |
| tokenizer.decode([token]), |
| float(candidates[token].logprob), |
| [(tokenizer.decode([i]), float(c.logprob)) for i, c in candidates.items()], |
| ) |
| for token, candidates in zip(completion.token_ids, completion.logprobs) |
| ] |
| probabilities, expected = core.score_probabilities(positions, score) |
| write_record(stream, { |
| "id": job["id"], |
| "task": job["task"], |
| "score": score, |
| "expected_score": round(expected, 4), |
| |
| "progress": round((expected - 1.0) / 4.0, 4), |
| "probabilities": {k: round(v, 4) for k, v in probabilities.items()}, |
| "rationale": rationale, |
| "duration_s": round(collage.duration_s, 3), |
| "sampled_frames": len(indices), |
| "output": completion.text, |
| }) |
| log(f"terminal: {min(start + len(batch), len(prepared))}/{len(prepared)} episodes") |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
