Add LaTeX OCR: multimodal env on the new Task API + dataset streaming

docs/source/_toctree.yml

@@ -131,6 +131,8 @@
     title: OpenCode
   - local: environments/sophistry_bench_sprint
     title: Sophistry Bench Sprint
+  - local: environments/latex_ocr
+    title: LaTeX OCR
   title: Environments
 - isExpanded: false
   sections:

docs/source/environments/latex_ocr.md

@@ -0,0 +1,79 @@
+<!-- openenv-source: latex_ocr_env -->
+# LaTeX OCR Environment
+
+A dataset-backed, single-step (bandit) RL environment for **image → LaTeX**
+transcription, served through OpenEnv.
+
+- **Task**: the agent is shown an image of a math/text expression and must
+  return its LaTeX source.
+- **Dataset**: tasks are served from a Hugging Face dataset (default
+  [`unsloth/LaTeX_OCR`](https://huggingface.co/datasets/unsloth/LaTeX_OCR):
+  `image` + `text` columns, `train`/`test` splits) via the OpenEnv **Task API**.
+- **Reward**: `[(1 - exact_weight) * (1 - CER) + exact_weight * exact_match] * length_factor`,
+  where `CER` is the normalized character edit distance over whitespace-stripped
+  LaTeX. With the default `exact_weight=0.4` (tunable via `LATEX_OCR_EXACT_WEIGHT`),
+  partial answers score in `[0, 0.6]`; only an exact match reaches `1.0`.
+  Computed **server-side** against the hidden ground truth — the agent never
+  sees the target on `reset`. Dense and smooth, for stable RL training.
+- **Length guard**: because the score is whitespace-insensitive, a policy could
+  otherwise emit the correct answer followed by unlimited whitespace and still
+  score `1.0` — a reward hack that, in RL, shows up as completions drifting to
+  the generation-length cap. `length_factor` decays the reward once the *raw*
+  prediction grows past `LATEX_OCR_OVERLONG_RATIO`× the target length (default
+  `4.0`, floor `LATEX_OCR_OVERLONG_FLOOR=80` chars); normal LaTeX spacing is well
+  within the allowance. Set `LATEX_OCR_OVERLONG_RATIO=0` to disable. The rubric
+  also strips code fences / `$`-delimiters from raw completions itself.
+
+## Episode
+
+```
+reset(split="test", index=0)  -> observation { image_base64, prompt, ... }   # target hidden
+step(LatexOCRAction(latex=…)) -> reward, done=True { target_latex, exact_match, char_error_rate }
+```
+
+## Task API
+
+| Endpoint | Purpose |
+|---|---|
+| `GET  /latex_ocr_env/splits` | list splits (`train`, `test`) |
+| `POST /latex_ocr_env/num_tasks` | row count for a split |
+| `POST /latex_ocr_env/tasks` | all task specs for a split |
+| `POST /latex_ocr_env/task` | one task by `{split, index}` |
+| `POST /latex_ocr_env/task_range` | slice `{split, start, stop}` |
+
+Client helpers: `env.list_splits()`, `env.num_tasks(split)`,
+`env.get_task(split, index)`, `env.get_task_range(split, start, stop)`.
+
+## Run locally
+
+```bash
+# From the repo root. LATEX_OCR_SPLITS=test avoids the 380MB train download.
+PYTHONPATH=src:envs/latex_ocr_env LATEX_OCR_SPLITS=test \
+  uv run --with datasets --with pillow --with fastapi --with uvicorn --with websockets \
+  uvicorn server.app:app --host 0.0.0.0 --port 8000
+```
+
+Then drive it (needs `HF_TOKEN` for the real VLM policy):
+
+```bash
+HF_TOKEN=hf_xxx PYTHONPATH=src \
+  uv run --with datasets --with pillow --with requests --with websockets --with openai \
+  python envs/latex_ocr_env/validate.py --split test --num 3 \
+  --model Qwen/Qwen2.5-VL-7B-Instruct
+```
+
+## Configuration (env vars)
+
+| Var | Default | Meaning |
+|---|---|---|
+| `LATEX_OCR_DATASET` | `unsloth/LaTeX_OCR` | source dataset |
+| `LATEX_OCR_IMAGE_COLUMN` | `image` | image column |
+| `LATEX_OCR_TEXT_COLUMN` | `text` | ground-truth LaTeX column |
+| `LATEX_OCR_SPLITS` | `train,test` | splits to expose |
+| `LATEX_OCR_MAX_ROWS` | — | cap rows per split (dev) |
+| `LATEX_OCR_EXACT_WEIGHT` | `0.4` | exact-match share of the reward |
+| `LATEX_OCR_OVERLONG_RATIO` | `4.0` | raw length allowed as a multiple of the target before the reward decays (`0` disables the length guard) |
+| `LATEX_OCR_OVERLONG_FLOOR` | `80` | minimum allowed raw length (chars) for short targets |
+
+Swap in any `(image, latex)` dataset by pointing `LATEX_OCR_DATASET` at it (and
+the column vars if they differ).

envs/latex_ocr_env/README.md

@@ -0,0 +1,88 @@
+---
+title: LaTeX OCR Env
+emoji: 📐
+colorFrom: blue
+colorTo: indigo
+sdk: docker
+app_port: 8000
+pinned: false
+---
+
+# LaTeX OCR Environment
+
+A dataset-backed, single-step (bandit) RL environment for **image → LaTeX**
+transcription, served through OpenEnv.
+
+- **Task**: the agent is shown an image of a math/text expression and must
+  return its LaTeX source.
+- **Dataset**: tasks are served from a Hugging Face dataset (default
+  [`unsloth/LaTeX_OCR`](https://huggingface.co/datasets/unsloth/LaTeX_OCR):
+  `image` + `text` columns, `train`/`test` splits) via the OpenEnv **Task API**.
+- **Reward**: `[(1 - exact_weight) * (1 - CER) + exact_weight * exact_match] * length_factor`,
+  where `CER` is the normalized character edit distance over whitespace-stripped
+  LaTeX. With the default `exact_weight=0.4` (tunable via `LATEX_OCR_EXACT_WEIGHT`),
+  partial answers score in `[0, 0.6]`; only an exact match reaches `1.0`.
+  Computed **server-side** against the hidden ground truth — the agent never
+  sees the target on `reset`. Dense and smooth, for stable RL training.
+- **Length guard**: because the score is whitespace-insensitive, a policy could
+  otherwise emit the correct answer followed by unlimited whitespace and still
+  score `1.0` — a reward hack that, in RL, shows up as completions drifting to
+  the generation-length cap. `length_factor` decays the reward once the *raw*
+  prediction grows past `LATEX_OCR_OVERLONG_RATIO`× the target length (default
+  `4.0`, floor `LATEX_OCR_OVERLONG_FLOOR=80` chars); normal LaTeX spacing is well
+  within the allowance. Set `LATEX_OCR_OVERLONG_RATIO=0` to disable. The rubric
+  also strips code fences / `$`-delimiters from raw completions itself.
+
+## Episode
+
+```
+reset(split="test", index=0)  -> observation { image_base64, prompt, ... }   # target hidden
+step(LatexOCRAction(latex=…)) -> reward, done=True { target_latex, exact_match, char_error_rate }
+```
+
+## Task API
+
+| Endpoint | Purpose |
+|---|---|
+| `GET  /latex_ocr_env/splits` | list splits (`train`, `test`) |
+| `POST /latex_ocr_env/num_tasks` | row count for a split |
+| `POST /latex_ocr_env/tasks` | all task specs for a split |
+| `POST /latex_ocr_env/task` | one task by `{split, index}` |
+| `POST /latex_ocr_env/task_range` | slice `{split, start, stop}` |
+
+Client helpers: `env.list_splits()`, `env.num_tasks(split)`,
+`env.get_task(split, index)`, `env.get_task_range(split, start, stop)`.
+
+## Run locally
+
+```bash
+# From the repo root. LATEX_OCR_SPLITS=test avoids the 380MB train download.
+PYTHONPATH=src:envs/latex_ocr_env LATEX_OCR_SPLITS=test \
+  uv run --with datasets --with pillow --with fastapi --with uvicorn --with websockets \
+  uvicorn server.app:app --host 0.0.0.0 --port 8000
+```
+
+Then drive it (needs `HF_TOKEN` for the real VLM policy):
+
+```bash
+HF_TOKEN=hf_xxx PYTHONPATH=src \
+  uv run --with datasets --with pillow --with requests --with websockets --with openai \
+  python envs/latex_ocr_env/validate.py --split test --num 3 \
+  --model Qwen/Qwen2.5-VL-7B-Instruct
+```
+
+## Configuration (env vars)
+
+| Var | Default | Meaning |
+|---|---|---|
+| `LATEX_OCR_DATASET` | `unsloth/LaTeX_OCR` | source dataset |
+| `LATEX_OCR_IMAGE_COLUMN` | `image` | image column |
+| `LATEX_OCR_TEXT_COLUMN` | `text` | ground-truth LaTeX column |
+| `LATEX_OCR_SPLITS` | `train,test` | splits to expose |
+| `LATEX_OCR_MAX_ROWS` | — | cap rows per split (dev) |
+| `LATEX_OCR_EXACT_WEIGHT` | `0.4` | exact-match share of the reward |
+| `LATEX_OCR_OVERLONG_RATIO` | `4.0` | raw length allowed as a multiple of the target before the reward decays (`0` disables the length guard) |
+| `LATEX_OCR_OVERLONG_FLOOR` | `80` | minimum allowed raw length (chars) for short targets |
+
+Swap in any `(image, latex)` dataset by pointing `LATEX_OCR_DATASET` at it (and
+the column vars if they differ).

envs/latex_ocr_env/__init__.py

@@ -0,0 +1,12 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""LaTeX OCR Environment - dataset-backed, single-step RL for image→LaTeX."""
+
+from .client import LatexOCREnv
+from .models import LatexOCRAction, LatexOCRObservation
+
+__all__ = ["LatexOCRAction", "LatexOCRObservation", "LatexOCREnv"]

envs/latex_ocr_env/client.py

@@ -0,0 +1,137 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""
+LaTeX OCR Environment client.
+
+Gym-style client (reset/step over WebSocket) plus thin HTTP helpers for the
+Task API so a trainer can enumerate and select dataset tasks.
+
+Example:
+    >>> with LatexOCREnv(base_url="http://localhost:8000") as env:
+    ...     print(env.list_splits())            # ["train", "test"]
+    ...     print(env.num_tasks("test"))        # 7595
+    ...     result = env.reset(split="test", index=0)
+    ...     img = result.observation.image_base64
+    ...     # ... run a VLM to produce `latex` ...
+    ...     result = env.step(LatexOCRAction(latex=latex))
+    ...     print(result.reward, result.observation.target_latex)
+"""
+
+from __future__ import annotations
+
+from typing import Any, Optional
+from urllib.parse import urljoin
+
+import requests
+
+try:
+    from openenv.core.client_types import StepResult
+    from openenv.core.env_client import EnvClient
+    from openenv.core.env_server.types import State
+except ImportError:
+    from core.client_types import StepResult
+    from core.env_client import EnvClient
+    from core.env_server.types import State
+
+from .models import LatexOCRAction, LatexOCRObservation
+
+ENV_NAME = "latex_ocr_env"
+
+
+class LatexOCREnv(EnvClient[LatexOCRAction, LatexOCRObservation, State]):
+    """Client for the LaTeX OCR environment."""
+
+    def reset(
+        self,
+        split: str = "train",
+        index: Optional[int] = None,
+        seed: Optional[int] = None,
+        **kwargs: Any,
+    ) -> StepResult[LatexOCRObservation]:
+        """Reset to a specific dataset task (or a random one if ``index`` is None)."""
+        payload: dict[str, Any] = {"split": split}
+        if index is not None:
+            payload["index"] = index
+        if seed is not None:
+            payload["seed"] = seed
+        payload.update(kwargs)
+        return super().reset(**payload)
+
+    # --- Gym-style (de)serialization required by EnvClient ---
+    def _step_payload(self, action: LatexOCRAction) -> dict[str, Any]:
+        return action.model_dump()
+
+    def _parse_result(self, data: dict[str, Any]) -> StepResult[LatexOCRObservation]:
+        obs_data = dict(data.get("observation", data))
+        # Core serialization lifts reward/done to the top level and strips them from
+        # the observation payload; merge them back so observation.reward/.done match
+        # StepResult (consistent with other env clients).
+        reward = data.get("reward", obs_data.get("reward"))
+        done = data.get("done", obs_data.get("done"))
+        obs_data["reward"] = reward
+        obs_data["done"] = done
+        obs = LatexOCRObservation(**obs_data)
+        base = dict(observation=obs, reward=reward, done=done)
+        # Newer core's StepResult carries `info`; older core does not.
+        try:
+            return StepResult(**base, info=data.get("info", {}))
+        except TypeError:
+            return StepResult(**base)
+
+    def _parse_state(self, data: dict[str, Any]) -> State:
+        return State(**data)
+
+    # ------------------------------------------------------------------ #
+    # Task API (HTTP)                                                     #
+    # ------------------------------------------------------------------ #
+    def _http_base(self) -> str:
+        # Newer core exposes ``_base_url`` (http); older core stores only
+        # ``_ws_url`` (ws://host/ws). Derive an http base that works for both.
+        base = getattr(self, "_base_url", None)
+        if not base:
+            ws = getattr(self, "_ws_url", None)
+            if not ws:
+                raise RuntimeError("Task API requires an HTTP base URL")
+            base = ws[:-3] if ws.endswith("/ws") else ws
+        base = base.replace("wss://", "https://").replace("ws://", "http://")
+        return base if base.endswith("/") else base + "/"
+
+    def list_splits(self) -> list[str]:
+        resp = requests.get(
+            urljoin(self._http_base(), f"{ENV_NAME}/splits"), timeout=30
+        )
+        resp.raise_for_status()
+        return [s["name"] for s in resp.json()]
+
+    def num_tasks(self, split: str) -> int:
+        resp = requests.post(
+            urljoin(self._http_base(), f"{ENV_NAME}/num_tasks"),
+            json={"split": split},
+            timeout=60,
+        )
+        resp.raise_for_status()
+        return int(resp.json()["num_tasks"])
+
+    def get_task(self, split: str, index: int) -> dict[str, Any]:
+        resp = requests.post(
+            urljoin(self._http_base(), f"{ENV_NAME}/task"),
+            json={"split": split, "index": index},
+            timeout=60,
+        )
+        resp.raise_for_status()
+        return resp.json()["task"]
+
+    def get_task_range(
+        self, split: str, start: int | None = None, stop: int | None = None
+    ) -> list[dict[str, Any]]:
+        resp = requests.post(
+            urljoin(self._http_base(), f"{ENV_NAME}/task_range"),
+            json={"split": split, "start": start, "stop": stop},
+            timeout=120,
+        )
+        resp.raise_for_status()
+        return resp.json()["tasks"]