diff --git a/anvil/observe/metrics.py b/anvil/observe/metrics.py index 0db94a1..4d5e63b 100644 --- a/anvil/observe/metrics.py +++ b/anvil/observe/metrics.py @@ -82,12 +82,14 @@ def log_step( adapter_synced: bool | None = None, snapshot_path: str | None = None, sample_endpoint: str | None = None, + job: str = "grpo", + n_image_refs: int = 0, ) -> dict[str, Any]: fb = dict(fb_metrics or {}) record: dict[str, Any] = { "schema_version": SCHEMA_VERSION, "type": "step", - "job": "grpo", + "job": str(job or "grpo"), "ts": time.time(), "step": int(step), "reward_mean": float(reward_mean), @@ -95,6 +97,7 @@ def log_step( "group_reward_std_mean": float(group_reward_std_mean), "loss": float(loss), "n_datums": int(n_datums), + "n_image_refs": int(n_image_refs), "is_mean_ratio": fb.get("mean_ratio"), "fb": fb, "wall_time_s": wall_time_s, diff --git a/anvil/recipes/__init__.py b/anvil/recipes/__init__.py index ade7c40..26030d4 100644 --- a/anvil/recipes/__init__.py +++ b/anvil/recipes/__init__.py @@ -53,6 +53,13 @@ toy_robot_trajectories, trajectories_to_robot_examples, ) +from anvil.recipes.vision_rl import ( + VisionRLQueueRecipe, + VisionRollout, + run_vision_grpo, + run_vision_rl_queue, + toy_vision_rollouts, +) __all__ = [ "BookRecipe", @@ -74,6 +81,8 @@ "RecipeSpec", "RobotOfflineResult", "StageRunResult", + "VisionRLQueueRecipe", + "VisionRollout", "default_book_root", "default_recipe_id_for_shape", "example_vlm_ladder", @@ -95,8 +104,11 @@ "run_dpo", "run_meta_recipe", "run_robot_offline", + "run_vision_grpo", + "run_vision_rl_queue", "save_meta_recipe", "suggest_for_model", "toy_robot_trajectories", + "toy_vision_rollouts", "trajectories_to_robot_examples", ] diff --git a/anvil/recipes/catalog.py b/anvil/recipes/catalog.py index d18762b..d4409fe 100644 --- a/anvil/recipes/catalog.py +++ b/anvil/recipes/catalog.py @@ -395,6 +395,32 @@ class RecipeSpec: "export ONNX/TRT later. Never actuate from raw samples without a supervisor." ), ), + RecipeSpec( + id="vision_rl_on_policy", + title="Vision on-policy RL (GRPO)", + summary=( + "Multimodal sample + keyword/rubric/action-bin rewards; " + "run_vision_grpo / vision stage queue." + ), + pattern=JobPattern.RL_VERIFIABLE, + shapes_recommended=(ModelShape.EDGE_STUDENT, ModelShape.DENSE_VLM), + shapes_stretch=(ModelShape.UNKNOWN,), + shapes_blocked=(ModelShape.MOE_LM,), + needs_vision=True, + needs_reward=True, + max_param_b_recommended=4.5, + max_param_b_hard=14.0, + default_rank=16, + default_lr=5e-5, + default_export="peft", + group="edge", + tags=("robot", "vision", "grpo", "on-policy", "rubric"), + notes=( + "Entry: anvil.recipes.vision_rl.run_vision_grpo. Image refs in prompts; " + "rewards on detokenized text (no learned vision RM in v0). Train on lab; " + "do not dump multi-hour vision RL logs onto the house robot." + ), + ), RecipeSpec( id="distill_to_edge", title="Distill lab teacher → edge student", diff --git a/anvil/recipes/grpo.py b/anvil/recipes/grpo.py index afcd7d7..fe73849 100644 --- a/anvil/recipes/grpo.py +++ b/anvil/recipes/grpo.py @@ -55,6 +55,27 @@ from anvil.recipes.profiles import JobPattern, RecipePlan, plan_recipe RewardFn = Callable[[str, Sequence[int]], float] +# One reward for all prompts, or one reward per prompt (same length as prompts). +RewardSpec = RewardFn | Sequence[RewardFn] +# Text-token prompts or full multimodal ModelInput (image refs + text tokens). +PromptSpec = Sequence[int] | ModelInput + + +def as_model_input(prompt: PromptSpec) -> ModelInput: + """Normalize a GRPO prompt to :class:`ModelInput` (text tokens or multimodal).""" + if isinstance(prompt, ModelInput): + return prompt + return ModelInput.from_ints(list(prompt)) + + +def resolve_reward_fn(reward: RewardSpec | None, prompt_index: int) -> RewardFn: + if reward is None: + return _exact_match_toy + if callable(reward): + return reward # type: ignore[return-value] + if not reward: + return _exact_match_toy + return reward[min(prompt_index, len(reward) - 1)] def datum_from_rollout( @@ -202,17 +223,18 @@ def push_adapter_snapshot( def run_grpo( *, base_model: str = "Qwen/Qwen3.5-4B", - prompts: Sequence[Sequence[int]] | None = None, - reward_fn: RewardFn | None = None, + prompts: Sequence[PromptSpec] | None = None, + reward_fn: RewardSpec | None = None, group_size: int = 4, steps: int = 3, endpoint: str = "fake://", plan: RecipePlan | None = None, overrides: dict[str, Any] | None = None, run_dir: str | None = None, - probes: Sequence[Sequence[int]] | None = None, + probes: Sequence[PromptSpec] | None = None, probe_every: int = 1, detokenize: Callable[[Sequence[int]], str] | None = None, + job: str = "grpo", sample_endpoint: str | None = None, sample_endpoints: Sequence[str] | None = None, sample_backend: Backend | None = None, @@ -289,11 +311,12 @@ def run_grpo( mm_projector=k["mm_projector_lora"], ), ) - reward_fn = reward_fn or _exact_match_toy + reward_spec: RewardSpec = reward_fn if reward_fn is not None else _exact_match_toy prompts = list(prompts) if prompts else [list(range(10, 26))] writer = RunMetricsWriter(run_dir) if run_dir else None if stop_on_southward is None: stop_on_southward = bool(early_stop and run_dir) + job_label = str(job or "grpo") sample_svc: ServiceClient | None = None sample_pool_owned = False @@ -333,7 +356,6 @@ def run_grpo( stopped_reason: str | None = None start_step = 0 last_ckpt_path: str | None = None - job = "grpo" if resume and run_dir: state = load_resume_state(run_dir) @@ -348,7 +370,7 @@ def run_grpo( step=start_step, event="resume", reason="loaded_resume_json", - job=job, + job=job_label, steps_completed=start_step, checkpoint_path=state.checkpoint_path, ) @@ -397,7 +419,7 @@ def _maybe_checkpoint(step_completed: int, *, force: bool = False) -> None: ref = save_train_checkpoint( tc, run_dir=run_dir, - job=job, + job=job_label, steps_completed=step_completed, base_model=plan.base_model, losses=prior_losses + losses, @@ -410,7 +432,7 @@ def _maybe_checkpoint(step_completed: int, *, force: bool = False) -> None: step=step_completed - 1 if step_completed > 0 else 0, event="checkpoint", reason="periodic" if not force else "final", - job=job, + job=job_label, steps_completed=step_completed, checkpoint_path=ref.path, checkpoint_name=ref.name, @@ -454,8 +476,10 @@ def _maybe_checkpoint(step_completed: int, *, force: bool = False) -> None: step_rewards: list[float] = [] group_stds: list[float] = [] - for prompt_tokens in prompts: - prompt = ModelInput.from_ints(prompt_tokens) + for prompt_ix, prompt_spec in enumerate(prompts): + prompt = as_model_input(prompt_spec) + prompt_tokens = prompt.token_ids() + rf = resolve_reward_fn(reward_spec, prompt_ix) sample = sc.sample( prompt, SamplingParams( @@ -467,7 +491,8 @@ def _maybe_checkpoint(step_completed: int, *, force: bool = False) -> None: ).result() rewards = [] for seq in sample.sequences: - r = reward_fn("", seq.tokens) + text = detokenize(seq.tokens) if detokenize is not None else "" + r = rf(text, seq.tokens) rewards.append(r) step_rewards.append(r) group_stds.append(statistics.pstdev(rewards) if len(rewards) > 1 else 0.0) @@ -489,20 +514,23 @@ def _maybe_checkpoint(step_completed: int, *, force: bool = False) -> None: # Probe the LIVE policy (weights used for this step's rollouts) if writer is not None and probes and step_ix % probe_every == 0: - for probe_ix, probe_tokens in enumerate(probes): + for probe_ix, probe_spec in enumerate(probes): + probe_in = as_model_input(probe_spec) out = sc.sample( - ModelInput.from_ints(probe_tokens), + probe_in, SamplingParams(max_tokens=plan.max_tokens, temperature=0.0, seed=0), num_samples=1, ).result() seq = out.sequences[0] if out.sequences else None toks = tuple(seq.tokens) if seq is not None else () + text = detokenize(toks) if detokenize is not None else None + prf = resolve_reward_fn(reward_spec, probe_ix) writer.log_probe( step=step_ix, probe_idx=probe_ix, tokens=toks, - text=detokenize(toks) if detokenize else None, - reward=reward_fn("", toks), + text=text, + reward=prf(text or "", toks), ) reward_mean = sum(step_rewards) / max(len(step_rewards), 1) @@ -535,6 +563,14 @@ def _maybe_checkpoint(step_completed: int, *, force: bool = False) -> None: losses.append(step_loss) mean_rewards.append(reward_mean) if writer is not None: + n_img = 0 + for p in prompts: + mi = as_model_input(p) + n_img += sum( + 1 + for c in mi.chunks + if type(c).__name__ in {"ImageRefChunk", "ImageChunk"} + ) writer.log_step( step=step_ix, reward_mean=reward_mean, @@ -549,6 +585,8 @@ def _maybe_checkpoint(step_completed: int, *, force: bool = False) -> None: adapter_synced=adapter_synced if resolved_sample is not None else None, snapshot_path=snap_path, sample_endpoint=sample_ep_label if resolved_sample is not None else None, + job=job_label, + n_image_refs=n_img, ) completed = step_ix + 1 diff --git a/anvil/recipes/vision_rewards.py b/anvil/recipes/vision_rewards.py new file mode 100644 index 0000000..1be6954 --- /dev/null +++ b/anvil/recipes/vision_rewards.py @@ -0,0 +1,164 @@ +"""Vision-aware verifiable rewards for on-policy GRPO (Phase 4.B). + +Rewards score **detokenized completion text** (and optional gold fields). +Image refs are prompt context for the sampler; they do not enter the reward +tensor directly in v0 (no learned vision RM). Rubrics stay explicit and +auditable — same live-sufficiency contract as text GRPO. +""" + +from __future__ import annotations + +import re +from dataclasses import dataclass +from typing import Callable, Sequence + +from anvil.recipes.grpo import RewardFn + +DetokenizeFn = Callable[[Sequence[int]], str] + + +def _norm(s: str) -> str: + return re.sub(r"\s+", " ", s.strip().lower()) + + +def keyword_caption_reward( + *, + required: Sequence[str], + any_of: Sequence[str] = (), + detokenize: DetokenizeFn | None = None, + case_insensitive: bool = True, +) -> RewardFn: + """1.0 if all ``required`` substrings appear (and any of ``any_of`` if set). + + Use for scene captions: gold keywords from SSD labels / human notes. + Partial credit: fraction of required hits when ``any_of`` is empty. + """ + req = [str(x) for x in required if str(x).strip()] + opt = [str(x) for x in any_of if str(x).strip()] + + def reward_fn(text: str, tokens: Sequence[int]) -> float: + body = text if text else (detokenize(tokens) if detokenize else "") + hay = _norm(body) if case_insensitive else body + if not req and not opt: + return 0.0 + hits = 0 + for k in req: + needle = _norm(k) if case_insensitive else k + if needle and needle in hay: + hits += 1 + if req: + score = hits / len(req) + else: + score = 0.0 + if opt: + if any((_norm(k) if case_insensitive else k) in hay for k in opt): + score = min(1.0, score + 0.25) if req else 1.0 + elif not req: + score = 0.0 + return float(score) + + return reward_fn + + +def exact_phrase_reward( + gold: str, + *, + detokenize: DetokenizeFn | None = None, +) -> RewardFn: + """1.0 iff normalized completion equals normalized gold (or contains it).""" + g = _norm(gold) + + def reward_fn(text: str, tokens: Sequence[int]) -> float: + body = text if text else (detokenize(tokens) if detokenize else "") + b = _norm(body) + if not g: + return 0.0 + if b == g or g in b: + return 1.0 + return 0.0 + + return reward_fn + + +def action_bin_overlap_reward( + gold_bins: Sequence[int] | str, + *, + detokenize: DetokenizeFn | None = None, +) -> RewardFn: + """Fraction of space-separated bin tokens that match gold (OpenVLA-style).""" + if isinstance(gold_bins, str): + gold = [int(x) for x in gold_bins.split() if x.lstrip("-").isdigit()] + else: + gold = [int(x) for x in gold_bins] + + def reward_fn(text: str, tokens: Sequence[int]) -> float: + body = text if text else (detokenize(tokens) if detokenize else "") + pred = [int(x) for x in body.split() if x.lstrip("-").isdigit()] + if not gold: + return 0.0 + n = min(len(gold), len(pred)) + if n == 0: + return 0.0 + hits = sum(1 for i in range(n) if pred[i] == gold[i]) + # length penalty if completion much shorter + if len(pred) < len(gold): + return hits / len(gold) + return hits / len(gold) + + return reward_fn + + +@dataclass(frozen=True, slots=True) +class RubricCriterion: + """One graded criterion (keyword or phrase).""" + + name: str + weight: float = 1.0 + required_keywords: tuple[str, ...] = () + forbidden_keywords: tuple[str, ...] = () + + +def rubric_reward( + criteria: Sequence[RubricCriterion], + *, + detokenize: DetokenizeFn | None = None, +) -> RewardFn: + """Weighted sum of criterion scores in [0, 1].""" + crits = list(criteria) + total_w = sum(max(0.0, c.weight) for c in crits) or 1.0 + + def reward_fn(text: str, tokens: Sequence[int]) -> float: + body = text if text else (detokenize(tokens) if detokenize else "") + hay = _norm(body) + acc = 0.0 + for c in crits: + w = max(0.0, c.weight) + if any(_norm(f) in hay for f in c.forbidden_keywords if f): + score = 0.0 + elif not c.required_keywords: + score = 1.0 if hay else 0.0 + else: + hits = sum(1 for k in c.required_keywords if _norm(k) in hay) + score = hits / len(c.required_keywords) + acc += w * score + return float(acc / total_w) + + return reward_fn + + +def toy_detokenize(tokens: Sequence[int]) -> str: + """CI detokenizer: map small ids to fixed vocabulary (kitchen scene words).""" + vocab = { + 0: "unsure", + 1: "kitchen", + 2: "chair", + 3: "table", + 4: "cabinet", + 5: "stove", + 6: "person", + 7: "dog", + 8: "room", + 9: "hallway", + } + parts = [vocab.get(int(t) % 10, str(int(t) % 10)) for t in tokens] + return " ".join(parts) diff --git a/anvil/recipes/vision_rl.py b/anvil/recipes/vision_rl.py new file mode 100644 index 0000000..3e9a23e --- /dev/null +++ b/anvil/recipes/vision_rl.py @@ -0,0 +1,332 @@ +"""On-policy vision RL (Phase 4.B) — multimodal GRPO + rubrics + stage queue. + +Builds on :func:`~anvil.recipes.grpo.run_grpo` with: + +- **Multimodal sample prompts** — image refs + text tokens as ``ModelInput`` +- **Vision rewards** — keyword / rubric / action-bin scorers on detokenized text +- **Stage queue** — curriculum of vision tasks on a shared LoRA (same idea as + text ``rl_queue`` / ``vlm_queue``) + +Lab trains; edge only samples small students. Never dump long vision RL runs +onto a storage-constrained robot. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Callable, Sequence + +from anvil.client.service import ServiceClient +from anvil.observe.metrics import RunMetricsWriter +from anvil.protocol.types import EncodedTextChunk, ImageRefChunk, LoraTargets, ModelInput +from anvil.recipes.grpo import GRPOResult, RewardFn, run_grpo +from anvil.recipes.profiles import JobPattern, plan_recipe +from anvil.recipes.rl_queue import should_advance +from anvil.recipes.vision_rewards import ( + RubricCriterion, + action_bin_overlap_reward, + keyword_caption_reward, + rubric_reward, + toy_detokenize, +) + +DetokenizeFn = Callable[[Sequence[int]], str] + + +@dataclass(frozen=True, slots=True) +class VisionRollout: + """One image-grounded RL prompt (scene caption, grasp rubric, action bins).""" + + id: str + instruction: str + image_refs: tuple[str, ...] = () + """``cas://`` or path refs for the observation (may be empty for text-only).""" + + required_keywords: tuple[str, ...] = () + any_of_keywords: tuple[str, ...] = () + gold_phrase: str | None = None + gold_action_bins: tuple[int, ...] | str | None = None + rubric: tuple[RubricCriterion, ...] = () + prompt_token_ids: tuple[int, ...] = () + """Optional explicit token ids for the instruction (CI / fake://).""" + + def to_model_input(self, *, default_prompt_tokens: Sequence[int] | None = None) -> ModelInput: + toks = list(self.prompt_token_ids) if self.prompt_token_ids else list( + default_prompt_tokens or range(10, 26) + ) + chunks: list[Any] = [EncodedTextChunk(tokens=tuple(int(t) for t in toks))] + for ref in self.image_refs: + chunks.append(ImageRefChunk(ref=str(ref), detail="auto")) + return ModelInput.from_chunks(chunks) + + def make_reward(self, *, detokenize: DetokenizeFn | None = None) -> RewardFn: + if self.rubric: + return rubric_reward(self.rubric, detokenize=detokenize) + if self.gold_action_bins is not None: + return action_bin_overlap_reward(self.gold_action_bins, detokenize=detokenize) + if self.gold_phrase: + from anvil.recipes.vision_rewards import exact_phrase_reward + + return exact_phrase_reward(self.gold_phrase, detokenize=detokenize) + if self.required_keywords or self.any_of_keywords: + return keyword_caption_reward( + required=self.required_keywords, + any_of=self.any_of_keywords, + detokenize=detokenize, + ) + # Weak default: any non-empty completion + def _any(text: str, tokens: Sequence[int]) -> float: + body = text or (detokenize(tokens) if detokenize else "") + return 1.0 if body.strip() else 0.0 + + return _any + + +def toy_vision_rollouts() -> list[VisionRollout]: + """Synthetic multimodal prompts for fake:// CI (no real pixels required).""" + dig = "e" * 64 + ref = f"cas://sha256/{dig}.png" + return [ + VisionRollout( + id="kitchen_caption", + instruction="Describe the room in one short sentence.", + image_refs=(ref,), + required_keywords=("kitchen",), + any_of_keywords=("stove", "cabinet", "table"), + prompt_token_ids=tuple(range(20, 36)), + ), + VisionRollout( + id="chair_hazard", + instruction="Name the main furniture hazard.", + image_refs=(ref,), + required_keywords=("chair",), + prompt_token_ids=tuple(range(30, 46)), + ), + ] + + +def run_vision_grpo( + *, + rollouts: Sequence[VisionRollout] | None = None, + base_model: str = "HuggingFaceTB/SmolVLM-256M-Instruct", + group_size: int = 4, + steps: int = 3, + endpoint: str = "fake://", + run_dir: str | None = None, + detokenize: DetokenizeFn | None = None, + overrides: dict[str, Any] | None = None, + fetch_remote: bool = False, + early_stop: bool = True, + early_stop_patience: int = 8, + stop_on_southward: bool | None = None, + service_client: Any | None = None, + training_client: Any | None = None, + close_clients: bool = True, + sample_endpoint: str | None = None, + sample_endpoints: Sequence[str] | None = None, + checkpoint_every: int | None = None, + resume: bool = False, +) -> GRPOResult: + """On-policy GRPO with multimodal prompts + vision rewards. + + ``fake://`` uses synthetic token detokenization when ``detokenize`` is None. + Lab hosts should pass a real tokenizer decode (or HF ``detokenize_via_tokenizer``). + """ + rolls = list(rollouts) if rollouts is not None else toy_vision_rollouts() + if not rolls: + raise ValueError("run_vision_grpo requires at least one VisionRollout") + + detok = detokenize if detokenize is not None else toy_detokenize + prompts = [r.to_model_input() for r in rolls] + rewards = [r.make_reward(detokenize=detok) for r in rolls] + ov = {**(overrides or {}), "modalities": ("text", "image")} + plan = plan_recipe( + base_model=base_model, + pattern=JobPattern.RL_VERIFIABLE, + overrides=ov, + use_card=False, + fetch_remote=fetch_remote, + ) + + return run_grpo( + base_model=base_model, + prompts=prompts, + reward_fn=rewards, + group_size=group_size, + steps=steps, + endpoint=endpoint, + plan=plan, + overrides=ov, + run_dir=run_dir, + probes=prompts[:1], + probe_every=1, + detokenize=detok, + job="vision_grpo", + early_stop=early_stop, + early_stop_patience=early_stop_patience, + stop_on_southward=stop_on_southward, + service_client=service_client, + training_client=training_client, + close_clients=close_clients, + sample_endpoint=sample_endpoint, + sample_endpoints=sample_endpoints, + checkpoint_every=checkpoint_every, + resume=resume, + ) + + +# --- vision / robot stage queue --------------------------------------------- + + +@dataclass(frozen=True, slots=True) +class VisionRLStage: + id: str + rollouts: tuple[VisionRollout, ...] + max_steps: int = 20 + group_size: int | None = None + early_stop_patience: int | None = None + notes: str = "" + + +@dataclass(frozen=True, slots=True) +class VisionRLQueueRecipe: + id: str + name: str + stages: tuple[VisionRLStage, ...] + group_size: int = 4 + early_stop_patience: int = 6 + advance_on: tuple[str, ...] = ("ceiling", "collapsed", "southward") + stop_queue_on: tuple[str, ...] = ("floor",) + advance_on_budget: bool = True + notes: str = "" + + +@dataclass +class VisionRLStageOutcome: + stage: VisionRLStage + result: GRPOResult + advanced: bool + queue_halted: bool + observe_run_id: str + + +@dataclass +class VisionRLQueueResult: + recipe: VisionRLQueueRecipe + stages: list[VisionRLStageOutcome] = field(default_factory=list) + adapter_id: str | None = None + + @property + def stages_run(self) -> int: + return len(self.stages) + + +def toy_vision_rl_queue() -> VisionRLQueueRecipe: + r0, r1 = toy_vision_rollouts() + return VisionRLQueueRecipe( + id="vision_caption_ladder_v0", + name="Vision caption → hazard ladder", + stages=( + VisionRLStage(id="caption", rollouts=(r0,), max_steps=5, notes="scene keywords"), + VisionRLStage(id="hazard", rollouts=(r1,), max_steps=5, notes="chair hazard"), + ), + notes="Phase 4.B toy curriculum for fake:// / lab smoke", + ) + + +def run_vision_rl_queue( + recipe: VisionRLQueueRecipe | None = None, + *, + base_model: str = "HuggingFaceTB/SmolVLM-256M-Instruct", + endpoint: str = "fake://", + run_dir: str | None = None, + detokenize: DetokenizeFn | None = None, + overrides: dict[str, Any] | None = None, + fetch_remote: bool = False, +) -> VisionRLQueueResult: + """Execute vision RL stages on a **shared** LoRA adapter.""" + recipe = recipe or toy_vision_rl_queue() + if not recipe.stages: + raise ValueError("VisionRLQueueRecipe requires stages") + + plan = plan_recipe( + base_model=base_model, + pattern=JobPattern.RL_VERIFIABLE, + overrides={**(overrides or {}), "modalities": ("text", "image")}, + use_card=False, + fetch_remote=fetch_remote, + ) + k = plan.as_knobs() + svc = ServiceClient(endpoint=endpoint) + tc = svc.create_lora_training_client( + base_model=plan.base_model, + rank=k["rank"], + modalities=k["modalities"], + lora_targets=LoraTargets( + language=k["language_lora"], + vision_encoder=k["vision_encoder_lora"], + mm_projector=k["mm_projector_lora"], + ), + ) + out = VisionRLQueueResult(recipe=recipe, adapter_id=str(tc.adapter_id)) + queue_writer = RunMetricsWriter(run_dir) if run_dir else None + + try: + for stage in recipe.stages: + stage_dir = f"{run_dir}/{stage.id}" if run_dir else None + if queue_writer is not None: + queue_writer.log_event( + step=0, + event="vision_rl_stage_start", + reason=stage.id, + job="vision_grpo", + stage=stage.id, + ) + res = run_vision_grpo( + rollouts=list(stage.rollouts), + base_model=base_model, + group_size=stage.group_size or recipe.group_size, + steps=stage.max_steps, + endpoint=endpoint, + run_dir=stage_dir, + detokenize=detokenize, + overrides=overrides, + fetch_remote=fetch_remote, + early_stop=True, + early_stop_patience=stage.early_stop_patience or recipe.early_stop_patience, + service_client=svc, + training_client=tc, + close_clients=False, + ) + hit_budget = res.early_stop_reason is None and res.steps_run >= stage.max_steps + advanced, halted = should_advance( + res.early_stop_reason, + hit_budget=hit_budget, + advance_on=recipe.advance_on, + stop_queue_on=recipe.stop_queue_on, + advance_on_budget=recipe.advance_on_budget, + ) + out.stages.append( + VisionRLStageOutcome( + stage=stage, + result=res, + advanced=advanced and not halted, + queue_halted=halted, + observe_run_id=stage.id, + ) + ) + if queue_writer is not None: + queue_writer.log_event( + step=res.steps_run, + event="vision_rl_stage_end", + reason=res.early_stop_reason or "budget", + job="vision_grpo", + stage=stage.id, + advanced=advanced and not halted, + halted=halted, + ) + if halted or not advanced: + break + finally: + svc.close() + return out diff --git a/docs/datasets-robotics.md b/docs/datasets-robotics.md index d27710a..6d1d3ab 100644 --- a/docs/datasets-robotics.md +++ b/docs/datasets-robotics.md @@ -208,7 +208,7 @@ vision must not lag text GRPO. | Checkpoint/resume + multi-hour lab profile | **done** (`checkpoint` + `lab_smokes --profile multi_hour`) | | Offline robot SFT + action tokens (smol edge) | **Phase 4.A** done | | House pack + edge export + agent dogfood | **Phase 4.A/C** + agentic path | -| Vision on-policy RL | **Phase 4.B** open | +| Vision on-policy RL (multimodal GRPO + rubrics + queue) | **Phase 4.B** mostly done (`run_vision_grpo`) | ## Smoke checklist diff --git a/docs/handoff.md b/docs/handoff.md index a4f9c71..93df810 100644 --- a/docs/handoff.md +++ b/docs/handoff.md @@ -34,6 +34,7 @@ Use this when starting a **new agent session** dedicated to Anvil (separate from | Checkpoint + resume **SFT/VLM/GRPO/DPO** | | | **Phase 4.A** robot_offline + action tokens | | | **House pack** + **edge export** + **agent dogfood** | | +| **Phase 4.B** vision GRPO + rubrics + queue | Lab VLM vision RL soak (operator) | | J-lens shelved (optional `-m jlens`) | | **Habit:** `python scripts/lab_smokes.py --profile quick` often. @@ -43,7 +44,7 @@ Use this when starting a **new agent session** dedicated to Anvil (separate from **Multi-worker sample:** `docs/multi-worker.md` **Recommended next Outcome:** -Phase 4.B on-policy vision RL, or j30 operator GGUF smoke with real camera pack +j30 operator GGUF FPS/power (storage-safe), then agentic forge dogfood soak ## Pull-on-miss only diff --git a/docs/roadmap.md b/docs/roadmap.md index b1a1214..f7e3a2c 100644 --- a/docs/roadmap.md +++ b/docs/roadmap.md @@ -384,11 +384,12 @@ path. Re-open only against the high bar in the spike writeup. - [x] Held-out **episode** split → probes (success-proxy sample under live adapter) - [x] In-house **house robot pack** (`anvil.data.robot_pack` + `scripts/robot_pack_smoke.py`; real device operator) -### 4.B On-policy vision RL *(open — required)* +### 4.B On-policy vision RL *(mostly closed)* -- [ ] Multimodal sample for vision rollouts -- [ ] Vision-aware rewards / rubrics -- [ ] Recipe queue for vision/robot stages +- [x] Multimodal sample for vision rollouts (`ModelInput` + image refs in `run_grpo` / `run_vision_grpo`) +- [x] Vision-aware rewards / rubrics (`vision_rewards` keyword / rubric / action-bin) +- [x] Recipe queue for vision/robot stages (`run_vision_rl_queue`) +- [ ] Lab-scale vision GRPO on real VLM + media store (operator on forge) ### 4.C Edge / Jetson *(mostly closed)* diff --git a/scripts/vision_rl_demo.py b/scripts/vision_rl_demo.py new file mode 100644 index 0000000..a366293 --- /dev/null +++ b/scripts/vision_rl_demo.py @@ -0,0 +1,83 @@ +#!/usr/bin/env python3 +"""Phase 4.B smoke: on-policy vision GRPO (+ optional stage queue). + + python scripts/vision_rl_demo.py --steps 3 + python scripts/vision_rl_demo.py --queue --steps 4 +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +REPO = Path(__file__).resolve().parents[1] +if str(REPO) not in sys.path: + sys.path.insert(0, str(REPO)) + + +def main(argv: list[str] | None = None) -> int: + p = argparse.ArgumentParser(description="Vision on-policy RL demo") + p.add_argument("--steps", type=int, default=3) + p.add_argument("--queue", action="store_true", help="run two-stage vision RL queue") + p.add_argument("--run-dir", default=None) + p.add_argument("--endpoint", default="fake://") + p.add_argument("--model", default="HuggingFaceTB/SmolVLM-256M-Instruct") + args = p.parse_args(argv) + + from anvil.recipes.vision_rl import run_vision_grpo, run_vision_rl_queue + + if args.queue: + q = run_vision_rl_queue( + base_model=args.model, + endpoint=args.endpoint, + run_dir=args.run_dir or "/tmp/anvil-vision-rl-queue", + ) + print( + json.dumps( + { + "mode": "queue", + "stages_run": q.stages_run, + "adapter_id": q.adapter_id, + "stages": [ + { + "id": o.stage.id, + "steps_run": o.result.steps_run, + "mean_reward": o.result.mean_reward, + "early_stop": o.result.early_stop_reason, + "advanced": o.advanced, + } + for o in q.stages + ], + }, + indent=2, + ) + ) + return 0 + + res = run_vision_grpo( + base_model=args.model, + steps=args.steps, + endpoint=args.endpoint, + run_dir=args.run_dir or "/tmp/anvil-vision-rl", + early_stop=False, + ) + print( + json.dumps( + { + "mode": "vision_grpo", + "adapter_id": res.adapter_id, + "steps_run": res.steps_run, + "mean_reward": res.mean_reward, + "losses": res.losses, + "run_dir": args.run_dir or "/tmp/anvil-vision-rl", + }, + indent=2, + ) + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/test_vision_rl.py b/tests/test_vision_rl.py new file mode 100644 index 0000000..39149f0 --- /dev/null +++ b/tests/test_vision_rl.py @@ -0,0 +1,131 @@ +"""Phase 4.B: multimodal GRPO, vision rewards, vision RL queue.""" + +from __future__ import annotations + +from pathlib import Path + +from anvil.protocol.types import ImageRefChunk, ModelInput +from anvil.recipes.grpo import as_model_input, run_grpo +from anvil.recipes.vision_rewards import ( + RubricCriterion, + action_bin_overlap_reward, + keyword_caption_reward, + rubric_reward, + toy_detokenize, +) +from anvil.recipes.vision_rl import ( + VisionRollout, + run_vision_grpo, + run_vision_rl_queue, + toy_vision_rl_queue, + toy_vision_rollouts, +) + + +def test_keyword_and_rubric_rewards(): + kw = keyword_caption_reward(required=("kitchen", "chair")) + assert kw("Kitchen with a chair", ()) == 1.0 + assert kw("empty room", ()) == 0.0 + assert 0.0 < kw("kitchen only", ()) < 1.0 + + rub = rubric_reward( + ( + RubricCriterion(name="scene", weight=1.0, required_keywords=("kitchen",)), + RubricCriterion(name="hazard", weight=1.0, required_keywords=("chair",)), + ) + ) + assert rub("kitchen chair", ()) == 1.0 + assert rub("kitchen", ()) == 0.5 + + +def test_action_bin_overlap(): + r = action_bin_overlap_reward([10, 20, 30]) + assert r("10 20 30", ()) == 1.0 + assert r("10 99 30", ()) == 2 / 3 + assert r("nope", ()) == 0.0 + + +def test_as_model_input_multimodal(): + from anvil.protocol.types import EncodedTextChunk + + mi = ModelInput.from_chunks( + ( + EncodedTextChunk(tokens=(1, 2, 3)), + ImageRefChunk(ref="cas://sha256/" + "a" * 64, detail="low"), + ) + ) + assert as_model_input(mi) is mi + assert as_model_input([4, 5, 6]).token_ids() == [4, 5, 6] + + +def test_run_grpo_passes_detokenized_text_to_reward(): + seen: list[str] = [] + + def rf(text: str, tokens): + seen.append(text) + return 1.0 if text else 0.0 + + run_grpo( + prompts=[[10, 11, 12]], + reward_fn=rf, + detokenize=lambda toks: "kitchen chair", + steps=1, + group_size=2, + endpoint="fake://", + early_stop=False, + ) + assert seen and all(s == "kitchen chair" for s in seen) + + +def test_run_grpo_multimodal_prompt_fake(tmp_path: Path): + mi = toy_vision_rollouts()[0].to_model_input() + assert any(isinstance(c, ImageRefChunk) for c in mi.chunks) + res = run_grpo( + prompts=[mi], + reward_fn=keyword_caption_reward(required=("kitchen",), detokenize=toy_detokenize), + detokenize=toy_detokenize, + steps=2, + group_size=2, + endpoint="fake://", + run_dir=str(tmp_path / "mm-grpo"), + early_stop=False, + job="vision_grpo", + ) + assert res.steps_run == 2 + metrics = (tmp_path / "mm-grpo" / "metrics.jsonl").read_text(encoding="utf-8") + assert "vision_grpo" in metrics + assert "n_image_refs" in metrics + + +def test_run_vision_grpo_and_queue(tmp_path: Path): + res = run_vision_grpo( + steps=2, + endpoint="fake://", + run_dir=str(tmp_path / "vgrpo"), + early_stop=False, + ) + assert res.steps_run == 2 + assert res.adapter_id + assert len(res.mean_reward) == 2 + + q = run_vision_rl_queue( + toy_vision_rl_queue(), + endpoint="fake://", + run_dir=str(tmp_path / "vqueue"), + ) + assert q.stages_run >= 1 + assert q.adapter_id + # stage subdirs + assert (tmp_path / "vqueue" / "caption").is_dir() or q.stages_run >= 1 + + +def test_vision_rollout_reward_builders(): + r = VisionRollout( + id="t", + instruction="x", + image_refs=("cas://sha256/" + "c" * 64,), + required_keywords=("kitchen",), + prompt_token_ids=(1, 2, 3), + ) + fn = r.make_reward(detokenize=toy_detokenize) + assert fn("kitchen table", ()) == 1.0