Initial commit

skillopt/envs/mmrb/__init__.py

@@ -0,0 +1,2 @@
+"""MMRB environment package."""
+

skillopt/envs/mmrb/adapter.py

@@ -0,0 +1,283 @@
+"""MMRB environment adapter for ReflACT."""
+from __future__ import annotations
+
+import json
+import os
+
+from skillopt.gradient.deep_probe import generate_deep_probe_instruction
+from skillopt.datasets.base import BatchSpec
+from skillopt.gradient.reflect import run_minibatch_reflect
+from skillopt.envs.base import EnvAdapter
+from skillopt.envs.mmrb.dataloader import MMRBDataLoader
+from skillopt.envs.mmrb.rollout import run_batch
+from skillopt.model import get_student_backend
+
+
+class MMRBAdapter(EnvAdapter):
+    """MMRB adapter."""
+
+    def build_reference_text(self, item: dict) -> str:
+        reasoning_steps = item.get("reasoning_steps") or []
+        if not reasoning_steps:
+            return ""
+
+        blocks: list[str] = []
+        for path_idx, path in enumerate(reasoning_steps, 1):
+            if not isinstance(path, list) or not path:
+                continue
+            lines = [f"### Reasoning Path {path_idx}"]
+            for step in path:
+                if not isinstance(step, dict):
+                    continue
+                step_no = step.get("reasoning step", "?")
+                step_type = str(step.get("reasoning type") or "").strip()
+                rationale = str(step.get("rationale") or "").strip()
+                if rationale:
+                    prefix = f"{step_no}. [{step_type}] " if step_type else f"{step_no}. "
+                    lines.append(prefix + rationale)
+            if len(lines) > 1:
+                blocks.append("\n".join(lines))
+        if not blocks:
+            return ""
+        return "## Reference Reasoning Steps\n" + "\n\n".join(blocks[:3])
+
+    def get_reference_metadata(self, item: dict) -> dict:
+        reasoning_steps = item.get("reasoning_steps") or []
+        path_count = 0
+        preview_parts: list[str] = []
+        for path in reasoning_steps:
+            if not isinstance(path, list) or not path:
+                continue
+            path_count += 1
+            first = path[0] if isinstance(path[0], dict) else {}
+            step_type = str(first.get("reasoning type") or "").strip()
+            rationale = str(first.get("rationale") or "").strip()
+            preview_parts.append(f"[path {path_count}] {step_type}: {rationale[:180]}")
+        if not path_count:
+            return {"fields": [], "preview": ""}
+        return {
+            "fields": ["reasoning_steps"],
+            "preview": "\n".join(preview_parts)[:500],
+        }
+
+    def __init__(
+        self,
+        split_dir: str = "",
+        data_path: str = "",
+        split_mode: str = "ratio",
+        split_ratio: str = "2:1:7",
+        split_seed: int = 42,
+        split_output_dir: str = "",
+        max_turns: int = 1,
+        workers: int = 16,
+        analyst_workers: int = 16,
+        failure_only: bool = False,
+        minibatch_size: int = 8,
+        edit_budget: int = 4,
+        seed: int = 42,
+        limit: int = 0,
+        image_detail: str = "auto",
+        use_deep_reflect: bool = False,
+        deep_reflect_failures: int = 4,
+        deep_reflect_successes: int = 2,
+    ) -> None:
+        self.max_turns = max_turns
+        self.workers = workers
+        self.analyst_workers = analyst_workers
+        self.failure_only = failure_only
+        self.minibatch_size = minibatch_size
+        self.edit_budget = edit_budget
+        self.image_detail = image_detail
+        self.use_deep_reflect = use_deep_reflect
+        self.deep_reflect_failures = deep_reflect_failures
+        self.deep_reflect_successes = deep_reflect_successes
+        self.dataloader = MMRBDataLoader(
+            split_dir=split_dir,
+            data_path=data_path,
+            split_mode=split_mode,
+            split_ratio=split_ratio,
+            split_seed=split_seed,
+            split_output_dir=split_output_dir,
+            seed=seed,
+            limit=limit,
+        )
+
+    def setup(self, cfg: dict) -> None:
+        super().setup(cfg)
+        self.dataloader.setup(cfg)
+
+    def get_dataloader(self):
+        return self.dataloader
+
+    def build_env_from_batch(self, batch: BatchSpec, **kwargs):
+        return list(batch.payload or [])
+
+    def build_train_env(self, batch_size: int, seed: int, **kwargs):
+        batch = self.dataloader.build_train_batch(batch_size=batch_size, seed=seed, **kwargs)
+        return self.build_env_from_batch(batch, **kwargs)
+
+    def build_eval_env(self, env_num: int, split: str, seed: int, **kwargs):
+        batch = self.dataloader.build_eval_batch(env_num=env_num, split=split, seed=seed, **kwargs)
+        return self.build_env_from_batch(batch, **kwargs)
+
+    def rollout(
+        self,
+        env_manager,
+        skill_content: str,
+        out_dir: str,
+        **kwargs,
+    ) -> list[dict]:
+        items: list[dict] = env_manager
+        return run_batch(
+            items=items,
+            out_root=out_dir,
+            skill_content=skill_content,
+            max_turns=self.max_turns,
+            workers=self.workers,
+            image_detail=self.image_detail,
+            diagnostic_mode=kwargs.get("diagnostic_mode", False),
+            diagnostic_instruction=kwargs.get("diagnostic_instruction", ""),
+            diagnostic_trace_context_by_id=kwargs.get("diagnostic_trace_context_by_id"),
+        )
+
+    def reflect(
+        self,
+        results: list[dict],
+        skill_content: str,
+        out_dir: str,
+        **kwargs,
+    ) -> list[dict | None]:
+        prediction_dir = kwargs.get("prediction_dir", os.path.join(out_dir, "predictions"))
+        patches_dir = kwargs.get("patches_dir", os.path.join(out_dir, "patches"))
+        random_seed = kwargs.get("random_seed")
+        step_buffer_context = kwargs.get("step_buffer_context", "")
+        meta_skill_context = kwargs.get("meta_skill_context", "")
+
+        return run_minibatch_reflect(
+            results=results,
+            skill_content=skill_content,
+            prediction_dir=prediction_dir,
+            patches_dir=patches_dir,
+            workers=self.analyst_workers,
+            failure_only=self.failure_only,
+            minibatch_size=self.minibatch_size,
+            edit_budget=self.edit_budget,
+            random_seed=random_seed,
+            error_system=self.get_error_minibatch_prompt(),
+            success_system=self.get_success_minibatch_prompt(),
+            step_buffer_context=step_buffer_context,
+            meta_skill_context=meta_skill_context,
+            update_mode=getattr(self, "_cfg", {}).get("skill_update_mode", "patch"),
+        )
+
+    def deep_reflect(
+        self,
+        results: list[dict],
+        skill_content: str,
+        out_dir: str,
+        **kwargs,
+    ) -> list[dict | None]:
+        if not self.use_deep_reflect:
+            return []
+
+        env_manager = kwargs.get("env_manager")
+        prediction_dir = kwargs.get("prediction_dir", os.path.join(out_dir, "predictions"))
+        random_seed = kwargs.get("random_seed")
+        step_buffer_context = kwargs.get("step_buffer_context", "")
+        meta_skill_context = kwargs.get("meta_skill_context", "")
+        codex_backend = get_student_backend() == "codex_exec"
+        selected_items = self.select_representative_items(
+            results,
+            env_manager if isinstance(env_manager, list) else None,
+            n_failures=self.deep_reflect_failures,
+            n_successes=self.deep_reflect_successes,
+            seed=random_seed,
+        )
+        if not selected_items:
+            return []
+        selected_ids = {str(item["id"]) for item in selected_items}
+        selected_results = [row for row in results if str(row.get("id")) in selected_ids]
+        selected_examples = self.attach_reference_context(selected_results, selected_items)
+        if codex_backend:
+            selected_examples = self.attach_codex_probe_context(selected_examples, prediction_dir)
+
+        reasoning_count = 0
+        selected_metadata = []
+        for item in selected_items:
+            meta = self.get_reference_metadata(item)
+            if meta["fields"]:
+                reasoning_count += 1
+            selected_metadata.append({
+                "id": str(item["id"]),
+                "task_type": str(item.get("subtask") or item.get("task_type") or "mmrb"),
+                "reference_fields": meta["fields"],
+                "reference_preview": meta["preview"],
+            })
+
+        deep_dir = os.path.join(out_dir, "deep_reflect")
+        rollout_dir = os.path.join(deep_dir, "rollout")
+        patches_dir = os.path.join(deep_dir, "patches")
+        os.makedirs(deep_dir, exist_ok=True)
+        print(
+            f"    [2b/6 DEEP REFLECT setup] selected={len(selected_items)} "
+            f"reference_fields=reasoning_steps({reasoning_count}/{len(selected_items)})"
+        )
+        probe = generate_deep_probe_instruction(
+            skill_content=skill_content,
+            items=selected_examples,
+            prediction_dir=prediction_dir,
+            system_prompt=self.get_codex_deep_probe_prompt() if codex_backend else self.get_deep_probe_prompt(),
+            step_buffer_context=step_buffer_context,
+            meta_skill_context=meta_skill_context,
+        )
+        if not probe:
+            return []
+        diagnostic_trace_context_by_id = None
+        if codex_backend:
+            selected_items, diagnostic_trace_context_by_id, probe = self.resolve_codex_probe_target(
+                selected_items=selected_items,
+                selected_examples=selected_examples,
+                prediction_dir=prediction_dir,
+                probe=probe,
+            )
+        probe_record = {
+            **probe,
+            "reference_summary": {
+                "selected_count": len(selected_items),
+                "field_counts": {"reasoning_steps": reasoning_count},
+            },
+            "selected_examples": selected_metadata,
+        }
+        with open(os.path.join(deep_dir, "probe.json"), "w", encoding="utf-8") as f:
+            json.dump(probe_record, f, ensure_ascii=False, indent=2)
+        deep_results = run_batch(
+            items=selected_items,
+            out_root=rollout_dir,
+            skill_content=skill_content,
+            max_turns=self.max_turns,
+            workers=min(self.workers, max(len(selected_items), 1)),
+            image_detail=self.image_detail,
+            diagnostic_mode=True,
+            diagnostic_instruction=probe["probe_instruction"],
+            diagnostic_trace_context_by_id=diagnostic_trace_context_by_id,
+        )
+        deep_results = self.attach_reference_context(deep_results, selected_items)
+        return run_minibatch_reflect(
+            results=deep_results,
+            skill_content=skill_content,
+            prediction_dir=os.path.join(rollout_dir, "predictions"),
+            patches_dir=patches_dir,
+            workers=self.analyst_workers,
+            failure_only=self.failure_only,
+            minibatch_size=self.minibatch_size,
+            edit_budget=self.edit_budget,
+            random_seed=random_seed,
+            error_system=self.get_error_minibatch_prompt(),
+            success_system=self.get_success_minibatch_prompt(),
+            step_buffer_context=step_buffer_context,
+            meta_skill_context=meta_skill_context,
+            update_mode=getattr(self, "_cfg", {}).get("skill_update_mode", "patch"),
+        )
+
+    def get_task_types(self) -> list[str]:
+        return self.dataloader.get_task_types()

skillopt/envs/mmrb/dataloader.py

@@ -0,0 +1,146 @@
+"""MMRB task dataloader."""
+from __future__ import annotations
+
+import glob
+import json
+import os
+import re
+from typing import Any
+
+from skillopt.datasets.base import SplitDataLoader
+
+
+# ── Raw data loading utilities (for preprocessing / standalone eval) ─────
+
+def _load_json(path: str) -> Any:
+    with open(path, encoding="utf-8") as f:
+        return json.load(f)
+
+
+def _iter_data_files(data_path: str) -> list[str]:
+    if not data_path:
+        return []
+    if os.path.isfile(data_path):
+        return [data_path]
+    if os.path.isdir(data_path):
+        nested = glob.glob(os.path.join(data_path, "**", "*_human.json"), recursive=True)
+        flat = glob.glob(os.path.join(data_path, "*_human.json"))
+        return sorted(set(nested + flat))
+    return []
+
+
+def _normalize_space(text: str) -> str:
+    return re.sub(r"\s+", " ", str(text or "").strip())
+
+
+def _normalize_item(item: dict, row_idx: int, source_path: str) -> dict | None:
+    question = _normalize_space(item.get("question") or "")
+    answer = _normalize_space(item.get("answer") or "")
+    raw_image_paths = item.get("image_paths") or []
+    if not question or not answer or not isinstance(raw_image_paths, list) or not raw_image_paths:
+        return None
+
+    base_dir = os.path.dirname(source_path)
+    image_paths: list[str] = []
+    for raw_path in raw_image_paths:
+        rel = str(raw_path or "").strip()
+        if not rel:
+            continue
+        abs_path = rel if os.path.isabs(rel) else os.path.abspath(os.path.join(base_dir, rel))
+        if os.path.exists(abs_path):
+            image_paths.append(abs_path)
+    if not image_paths:
+        return None
+
+    options_raw = item.get("options") or []
+    options = [_normalize_space(opt) for opt in options_raw if _normalize_space(opt)]
+    source = _normalize_space(item.get("source") or "unknown")
+    subtask = _normalize_space(item.get("subtask") or "unknown")
+    item_index = item.get("index", row_idx)
+    item_id = f"{source}:{subtask}:{item_index}"
+
+    return {
+        "id": item_id,
+        "source": source,
+        "subtask": subtask,
+        "task_type": subtask,
+        "question": question,
+        "answer": answer,
+        "options": options,
+        "is_choice": bool(options),
+        "image_paths": image_paths,
+        "reasoning_steps": item.get("reasoning_steps") or [],
+        "annotation_time": item.get("annotation_time"),
+        "source_path": os.path.abspath(source_path),
+    }
+
+
+def load_items(data_path: str) -> list[dict]:
+    """Load and normalise MMRB items from JSON files."""
+    files = _iter_data_files(data_path)
+    if not files:
+        raise ValueError(
+            "MMRB requires data_path to be a *_human.json file or a directory "
+            "containing extracted MMRB subtask folders."
+        )
+
+    items: list[dict] = []
+    for path in files:
+        raw = _load_json(path)
+        if not isinstance(raw, list):
+            raise ValueError(f"Expected JSON array in {path}, got {type(raw).__name__}")
+        for row_idx, item in enumerate(raw):
+            if not isinstance(item, dict):
+                continue
+            norm = _normalize_item(item, row_idx=row_idx, source_path=path)
+            if norm is not None:
+                items.append(norm)
+
+    if not items:
+        raise ValueError(f"No valid MMRB items loaded from {data_path}")
+    return items
+
+
+# ── Dataloader ───────────────────────────────────────────────────────────
+
+class MMRBDataLoader(SplitDataLoader):
+    """MMRB dataloader."""
+
+    def __init__(
+        self,
+        split_dir: str = "",
+        data_path: str = "",
+        split_mode: str = "ratio",
+        split_ratio: str = "2:1:7",
+        split_seed: int = 42,
+        split_output_dir: str = "",
+        seed: int = 42,
+        limit: int = 0,
+        **kwargs,
+    ) -> None:
+        super().__init__(
+            split_dir=split_dir,
+            data_path=data_path,
+            split_mode=split_mode,
+            split_ratio=split_ratio,
+            split_seed=split_seed,
+            split_output_dir=split_output_dir,
+            seed=seed,
+            limit=limit,
+        )
+        self._task_types: list[str] = []
+
+    def load_raw_items(self, data_path: str) -> list[dict]:
+        return load_items(data_path)
+
+    def setup(self, cfg: dict) -> None:
+        super().setup(cfg)
+        all_items = self.train_items + self.val_items + self.test_items
+        task_types = {
+            item.get("subtask") or item.get("task_type") or "unknown"
+            for item in all_items
+        }
+        self._task_types = sorted(task_types)
+
+    def get_task_types(self) -> list[str]:
+        return list(self._task_types)

skillopt/envs/mmrb/evaluator.py

@@ -0,0 +1,102 @@
+"""MMRB evaluation helpers."""
+from __future__ import annotations
+
+import re
+import string
+
+
+_EVAL_MODE = "mmrb_exact_match_v1"
+
+
+def normalize_text(text: str) -> str:
+    text = str(text or "").strip().lower()
+    text = "".join(ch for ch in text if ch not in string.punctuation)
+    return " ".join(text.split())
+
+
+def extract_answer(text: str | None) -> str:
+    raw = str(text or "").strip()
+    if not raw:
+        return ""
+
+    answer_tags = re.findall(r"<answer>\s*(.*?)\s*</answer>", raw, re.IGNORECASE | re.DOTALL)
+    if answer_tags:
+        return answer_tags[-1].strip()
+
+    bracket = re.findall(r"Answer\s*\[\s*(.*?)\s*\]", raw, re.IGNORECASE | re.DOTALL)
+    if bracket:
+        return bracket[-1].strip()
+
+    boxed = re.findall(r"\\boxed\{(.*?)\}", raw, re.IGNORECASE | re.DOTALL)
+    if boxed:
+        return boxed[-1].strip()
+
+    single = raw.strip().rstrip(".):")
+    if re.fullmatch(r"[A-Z]", single, re.IGNORECASE):
+        return single.strip()
+
+    patterns = [
+        r"final answer\s*(?:is)?\s*[:：]?\s*(.+)",
+        r"the answer is\s*[:：]?\s*(.+)",
+        r"answer\s*[:：]?\s*(.+)$",
+    ]
+    for pattern in patterns:
+        match = re.search(pattern, raw, re.IGNORECASE)
+        if match:
+            return match.group(1).strip().strip("*")
+
+    return raw
+
+
+def evaluate_item(*, item: dict, prediction_text: str) -> dict:
+    predicted_answer = extract_answer(prediction_text)
+    gold_answer = str(item.get("answer") or "").strip()
+    predicted_norm = normalize_text(predicted_answer)
+    gold_norm = normalize_text(gold_answer)
+
+    hard = 0.0
+    matched_gold = ""
+    predicted_label = ""
+    predicted_text = predicted_answer
+
+    if item.get("is_choice"):
+        predicted_label = str(predicted_answer).strip().upper().rstrip(".):")
+        if predicted_label == str(gold_answer).strip().upper():
+            hard = 1.0
+            matched_gold = gold_answer
+        else:
+            for option in item.get("options") or []:
+                label_match = re.match(r"\(?([A-Z])\)", option)
+                if not label_match:
+                    continue
+                label = label_match.group(1).upper()
+                option_text = option[label_match.end():].strip(" .:-")
+                if predicted_norm and normalize_text(option_text) == predicted_norm:
+                    predicted_label = label
+                    predicted_text = option_text
+                    break
+            if predicted_label == str(gold_answer).strip().upper():
+                hard = 1.0
+                matched_gold = gold_answer
+    else:
+        if predicted_norm and gold_norm and (
+            predicted_norm == gold_norm or predicted_norm in gold_norm or gold_norm in predicted_norm
+        ):
+            hard = 1.0
+            matched_gold = gold_answer
+
+    return {
+        "evaluation_mode": _EVAL_MODE,
+        "predicted_answer": predicted_answer,
+        "predicted_label": predicted_label,
+        "predicted_text": predicted_text,
+        "em": hard,
+        "f1": hard,
+        "sub_em": hard,
+        "matched_gold": matched_gold,
+    }
+
+
+def evaluation_mode() -> str:
+    return _EVAL_MODE
+

skillopt/envs/mmrb/prompts/rollout_system.md

@@ -0,0 +1,10 @@
+You are an expert multi-image reasoning agent.
+
+{skill_section}## Task Format
+You will receive a question grounded in multiple images.
+Use the image order exactly as presented in the prompt and compare evidence across images carefully.
+
+## Answer Format
+- Put the final answer inside <answer>...</answer>.
+- For multiple-choice questions, output only the single option letter inside <answer>...</answer>.
+- For open questions, output only the short final answer inside <answer>...</answer>.

skillopt/envs/mmrb/rollout.py

@@ -0,0 +1,439 @@
+"""MMRB rollout."""
+from __future__ import annotations
+
+import base64
+import json
+import mimetypes
+import os
+import re
+from concurrent.futures import ThreadPoolExecutor, as_completed
+
+from skillopt.envs.mmrb.evaluator import evaluate_item, evaluation_mode
+from skillopt.model import chat_student_messages, get_student_backend, is_student_exec_backend
+from skillopt.model.codex_harness import prepare_workspace, render_skill_md, run_student_exec
+from skillopt.prompts import load_prompt
+
+_IMAGE_REF_RE = re.compile(r"\{image#(\d+)\}", re.IGNORECASE)
+
+
+def _build_system(skill_content: str) -> str:
+    if skill_content.strip():
+        skill_section = f"## Skill\n{skill_content.strip()}\n\n"
+    else:
+        skill_section = ""
+    return load_prompt("rollout_system", env="mmrb").format(skill_section=skill_section)
+
+
+def _image_to_data_uri(path: str) -> str:
+    mime = mimetypes.guess_type(path)[0] or "image/png"
+    with open(path, "rb") as f:
+        encoded = base64.b64encode(f.read()).decode("ascii")
+    return f"data:{mime};base64,{encoded}"
+
+
+def _build_user_content(
+    item: dict,
+    image_detail: str,
+    *,
+    diagnostic_mode: bool = False,
+    diagnostic_instruction: str = "",
+    diagnostic_trace_context: str = "",
+) -> tuple[list[dict], str]:
+    raw_question = str(item["question"])
+    content: list[dict] = []
+    text_parts: list[str] = []
+    used_indices: set[int] = set()
+    cursor = 0
+
+    if diagnostic_trace_context.strip():
+        prefix = (
+            "## Previous Codex Trace Snapshot\n"
+            "This is a partial transcript from an earlier attempt. Use it as your current reasoning context.\n\n"
+            f"{diagnostic_trace_context.strip()}\n\n"
+        )
+        content.append({"type": "text", "text": prefix})
+        text_parts.append(prefix)
+
+    for match in _IMAGE_REF_RE.finditer(raw_question):
+        if match.start() > cursor:
+            chunk = raw_question[cursor:match.start()]
+            if chunk:
+                content.append({"type": "text", "text": chunk})
+                text_parts.append(chunk)
+
+        image_idx = int(match.group(1)) - 1
+        marker = f"[Image #{image_idx + 1}]"
+        text_parts.append(marker)
+        if 0 <= image_idx < len(item["image_paths"]):
+            image_url = {"url": _image_to_data_uri(item["image_paths"][image_idx])}
+            if image_detail and image_detail != "auto":
+                image_url["detail"] = image_detail
+            content.append({"type": "image_url", "image_url": image_url})
+            used_indices.add(image_idx)
+        else:
+            content.append({"type": "text", "text": marker})
+        cursor = match.end()
+
+    if cursor < len(raw_question):
+        tail = raw_question[cursor:]
+        if tail:
+            content.append({"type": "text", "text": tail})
+            text_parts.append(tail)
+
+    for idx, path in enumerate(item["image_paths"]):
+        if idx in used_indices:
+            continue
+        marker = f"\n[Additional Image #{idx + 1}]"
+        text_parts.append(marker)
+        content.append({"type": "text", "text": marker})
+        image_url = {"url": _image_to_data_uri(path)}
+        if image_detail and image_detail != "auto":
+            image_url["detail"] = image_detail
+        content.append({"type": "image_url", "image_url": image_url})
+
+    answer_instruction = (
+        "\n\nAnswer with the single correct option letter inside <answer>...</answer>."
+        if item.get("is_choice")
+        else "\n\nAnswer with the short final answer inside <answer>...</answer>."
+    )
+    content.append({"type": "text", "text": answer_instruction})
+    text_parts.append(answer_instruction)
+
+    if diagnostic_mode and diagnostic_instruction.strip():
+        diag_block = f"\n\n## Training Readout\n{diagnostic_instruction.strip()}"
+        content.append({"type": "text", "text": diag_block})
+        text_parts.append(diag_block)
+
+    return content, "".join(text_parts)
+
+
+def _build_messages(
+    item: dict,
+    skill_content: str,
+    image_detail: str,
+    *,
+    diagnostic_mode: bool = False,
+    diagnostic_instruction: str = "",
+) -> tuple[list[dict], str, str]:
+    system = _build_system(skill_content)
+    user_content, user_text = _build_user_content(
+        item,
+        image_detail,
+        diagnostic_mode=diagnostic_mode,
+        diagnostic_instruction=diagnostic_instruction,
+    )
+    messages = [
+        {"role": "system", "content": system},
+        {"role": "user", "content": user_content},
+    ]
+    return messages, system, user_text
+
+
+def _build_codex_skill(skill_content: str) -> str:
+    return render_skill_md(
+        skill_content,
+        description="Dynamic ReflACT skill for solving the current MMRB multi-image reasoning question.",
+        preamble=(
+            "Use this skill when solving the current multi-image reasoning task.\n"
+            "Inspect all attached images carefully and return the final answer inside <answer>...</answer>."
+        ),
+    )
+
+
+def _run_codex_once(
+    *,
+    pred_dir: str,
+    item: dict,
+    skill_content: str,
+    model: str,
+    timeout: int,
+    image_detail: str,
+    diagnostic_mode: bool = False,
+    diagnostic_instruction: str = "",
+    diagnostic_trace_context: str = "",
+    previous_response: str = "",
+) -> tuple[str, str, str, str]:
+    user_text = _build_user_content(
+        item,
+        image_detail,
+        diagnostic_mode=diagnostic_mode,
+        diagnostic_instruction=diagnostic_instruction,
+        diagnostic_trace_context=diagnostic_trace_context,
+    )[1]
+    task_parts = [user_text]
+    if previous_response:
+        task_parts.append(
+            "## Previous Attempt\n"
+            f"{previous_response}\n\n"
+            "Review the same images carefully and answer again."
+        )
+    task_text = "\n\n".join(task_parts)
+    skill_md = _build_codex_skill(skill_content)
+    work_dir = os.path.join(pred_dir, "codex_exec")
+    prepare_workspace(
+        work_dir=work_dir,
+        skill_md=skill_md,
+        task_text=task_text,
+        images=item["image_paths"],
+    )
+    prompt = (
+        "Use the `skillopt-student` skill available in this workspace.\n"
+        "Read `task.md`, inspect all attached images, and answer the question.\n"
+        "Keep the final answer inside <answer>...</answer>."
+    )
+    final_message, raw = run_student_exec(
+        work_dir=work_dir,
+        prompt=prompt,
+        model=model,
+        timeout=timeout,
+        images=item["image_paths"],
+    )
+    return final_message or raw, raw, skill_md, task_text
+
+
+def process_one(
+    item: dict,
+    out_root: str,
+    skill_content: str,
+    *,
+    max_turns: int = 1,
+    image_detail: str = "auto",
+    diagnostic_mode: bool = False,
+    diagnostic_instruction: str = "",
+    diagnostic_trace_context: str = "",
+) -> dict:
+    item_id = str(item["id"])
+    result = {
+        "id": item_id,
+        "question": item["question"],
+        "task_type": item.get("subtask") or item.get("task_type") or "mmrb",
+        "task_description": item["question"],
+        "hard": 0,
+        "soft": 0.0,
+        "predicted_answer": "",
+        "predicted_label": "",
+        "predicted_text": "",
+        "response": "",
+        "fail_reason": "",
+        "agent_ok": False,
+        "n_turns": 0,
+        "image_paths": item["image_paths"],
+        "gold_answer": item["answer"],
+        "evaluation_mode": evaluation_mode(),
+    }
+
+    try:
+        pred_dir = os.path.join(out_root, "predictions", item_id)
+        os.makedirs(pred_dir, exist_ok=True)
+
+        if is_student_exec_backend():
+            from skillopt.model import azure_openai as _llm
+
+            response = ""
+            conversation: list[dict] = [
+                {
+                    "role": "user",
+                    "content": item["question"] + "\n\n" + "\n".join(
+                        f"[image] {os.path.basename(path)}" for path in item["image_paths"]
+                    ),
+                }
+            ]
+            system_prompt = ""
+            user_text = ""
+            for turn in range(max_turns):
+                response, raw, system_prompt, user_text = _run_codex_once(
+                    pred_dir=pred_dir,
+                    item=item,
+                    skill_content=skill_content,
+                    model=_llm.STUDENT_DEPLOYMENT,
+                    timeout=120,
+                    image_detail=image_detail,
+                    diagnostic_mode=diagnostic_mode if turn == 0 else False,
+                    diagnostic_instruction=diagnostic_instruction if turn == 0 else "",
+                    diagnostic_trace_context=diagnostic_trace_context if turn == 0 else "",
+                    previous_response=response if turn > 0 else "",
+                )
+                conversation.append({"type": "message", "turn": turn + 1, "content": response})
+                if "<answer>" in response.lower():
+                    break
+
+            result["response"] = response
+            result["agent_ok"] = True
+            result["n_turns"] = len(conversation) - 1
+            with open(os.path.join(pred_dir, "student_system_prompt.txt"), "w", encoding="utf-8") as f:
+                f.write(system_prompt)
+            with open(os.path.join(pred_dir, "student_user_prompt.txt"), "w", encoding="utf-8") as f:
+                f.write(user_text)
+
+            eval_result = evaluate_item(item=item, prediction_text=response)
+            result["evaluation_mode"] = eval_result["evaluation_mode"]
+            result["predicted_answer"] = eval_result["predicted_answer"]
+            result["predicted_label"] = eval_result["predicted_label"]
+            result["predicted_text"] = eval_result["predicted_text"]
+            result["matched_gold"] = eval_result["matched_gold"]
+            result["hard"] = int(eval_result["em"])
+            result["soft"] = eval_result["f1"]
+            if not result["hard"]:
+                result["fail_reason"] = (
+                    f"predicted '{eval_result['predicted_answer']}' but expected '{item['answer']}'"
+                )
+            eval_detail = (
+                "[EVALUATION RESULT]\n"
+                f"Question: {item['question']}\n"
+                f"Predicted answer: {eval_result['predicted_answer']!r}\n"
+                f"Predicted label: {eval_result['predicted_label']!r}\n"
+                f"Gold answer: {item['answer']!r}\n"
+                f"Correct: {eval_result['em']}\n"
+            )
+            conversation.append({"role": "system", "content": eval_detail})
+            with open(os.path.join(pred_dir, "conversation.json"), "w", encoding="utf-8") as f:
+                json.dump(conversation, f, ensure_ascii=False, indent=2)
+            return result
+
+        messages, system_prompt, user_text = _build_messages(
+            item,
+            skill_content,
+            image_detail,
+            diagnostic_mode=diagnostic_mode,
+            diagnostic_instruction=diagnostic_instruction,
+            diagnostic_trace_context=diagnostic_trace_context,
+        )
+        response = ""
+        conversation: list[dict] = [
+            {
+                "role": "user",
+                "content": user_text + "\n\n" + "\n".join(
+                    f"[image] {os.path.basename(path)}" for path in item["image_paths"]
+                ),
+            }
+        ]
+
+        for turn in range(max_turns):
+            if turn == 0:
+                resp_text, _ = chat_student_messages(
+                    messages=messages,
+                    max_completion_tokens=768,
+                    retries=5,
+                    stage="rollout",
+                )
+            else:
+                refinement_messages = [
+                    messages[0],
+                    messages[1],
+                    {"role": "assistant", "content": response},
+                    {
+                        "role": "user",
+                        "content": "Review the same images carefully and answer again. Keep the final answer inside <answer>...</answer>.",
+                    },
+                ]
+                resp_text, _ = chat_student_messages(
+                    messages=refinement_messages,
+                    max_completion_tokens=512,
+                    retries=5,
+                    stage="rollout",
+                )
+            response = resp_text
+            conversation.append({"type": "message", "turn": turn + 1, "content": resp_text})
+            if "<answer>" in resp_text.lower():
+                break
+
+        result["response"] = response
+        result["agent_ok"] = True
+        result["n_turns"] = len(conversation) - 1
+
+        with open(os.path.join(pred_dir, "student_system_prompt.txt"), "w", encoding="utf-8") as f:
+            f.write(system_prompt)
+        with open(os.path.join(pred_dir, "student_user_prompt.txt"), "w", encoding="utf-8") as f:
+            f.write(user_text)
+
+        eval_result = evaluate_item(item=item, prediction_text=response)
+        result["evaluation_mode"] = eval_result["evaluation_mode"]
+        result["predicted_answer"] = eval_result["predicted_answer"]
+        result["predicted_label"] = eval_result["predicted_label"]
+        result["predicted_text"] = eval_result["predicted_text"]
+        result["matched_gold"] = eval_result["matched_gold"]
+        result["hard"] = int(eval_result["em"])
+        result["soft"] = eval_result["f1"]
+        if not result["hard"]:
+            result["fail_reason"] = (
+                f"predicted '{eval_result['predicted_answer']}' but expected '{item['answer']}'"
+            )
+
+        eval_detail = (
+            "[EVALUATION RESULT]\n"
+            f"Question: {item['question']}\n"
+            f"Predicted answer: {eval_result['predicted_answer']!r}\n"
+            f"Predicted label: {eval_result['predicted_label']!r}\n"
+            f"Gold answer: {item['answer']!r}\n"
+            f"Correct: {eval_result['em']}\n"
+        )
+        conversation.append({"role": "system", "content": eval_detail})
+        with open(os.path.join(pred_dir, "conversation.json"), "w", encoding="utf-8") as f:
+            json.dump(conversation, f, ensure_ascii=False, indent=2)
+    except Exception as e:  # noqa: BLE001
+        result["fail_reason"] = f"error: {e}"
+    return result
+
+
+def run_batch(
+    items: list[dict],
+    out_root: str,
+    skill_content: str,
+    *,
+    max_turns: int = 1,
+    workers: int = 16,
+    image_detail: str = "auto",
+    diagnostic_mode: bool = False,
+    diagnostic_instruction: str = "",
+    diagnostic_trace_context_by_id: dict[str, str] | None = None,
+) -> list[dict]:
+    results_path = os.path.join(out_root, "results.jsonl")
+    os.makedirs(out_root, exist_ok=True)
+
+    expected_eval_mode = evaluation_mode()
+    done_ids: set[str] = set()
+    existing: list[dict] = []
+    rewrite_results = False
+    if os.path.exists(results_path):
+        with open(results_path, encoding="utf-8") as f:
+            for line in f:
+                try:
+                    row = json.loads(line)
+                    if row.get("evaluation_mode") != expected_eval_mode:
+                        rewrite_results = True
+                        continue
+                    done_ids.add(str(row["id"]))
+                    existing.append(row)
+                except Exception:
+                    rewrite_results = True
+
+    pending = [item for item in items if str(item["id"]) not in done_ids]
+    if not pending and not rewrite_results:
+        return existing
+
+    results = list(existing)
+    file_mode = "w" if rewrite_results else "a"
+    with open(results_path, file_mode, encoding="utf-8") as outf, ThreadPoolExecutor(max_workers=workers) as ex:
+        if rewrite_results:
+            for row in existing:
+                outf.write(json.dumps(row, ensure_ascii=False) + "\n")
+        futs = {
+            ex.submit(
+                process_one,
+                item,
+                out_root,
+                skill_content,
+                max_turns=max_turns,
+                image_detail=image_detail,
+                diagnostic_mode=diagnostic_mode,
+                diagnostic_instruction=diagnostic_instruction,
+                diagnostic_trace_context=(diagnostic_trace_context_by_id or {}).get(str(item["id"]), ""),
+            ): item
+            for item in pending
+        }
+        for fut in as_completed(futs):
+            row = fut.result()
+            results.append(row)
+            outf.write(json.dumps(row, ensure_ascii=False) + "\n")
+            outf.flush()
+    return results

skillopt/envs/mmrb/skills/initial.md

@@ -0,0 +1,17 @@
+# MMRB Multi-Image Reasoning Heuristics
+
+## Cross-Image Alignment
+- Track the role of each image by its index and compare evidence across all referenced images before deciding.
+- When the question depends on sequence, correspondence, or retrieval, verify the relation between images instead of judging each image independently.
+
+## Option Elimination
+- For multiple-choice tasks, compare all options and reject choices that match only part of the visual evidence.
+- If options differ by a small visual detail, use the most discriminative cue rather than a coarse scene impression.
+
+## Open Answers
+- For open-ended tasks, give the shortest answer that is fully supported by the combined images.
+- Preserve exact entities, attributes, counts, and directions when the images support them directly.
+
+## Final Answer
+- Output only the final answer inside <answer>...</answer>.
+