Add LaTeX OCR: multimodal env on the new Task API + dataset streaming - #1003
Open
adithya-s-k wants to merge 26 commits into
Open
Add LaTeX OCR: multimodal env on the new Task API + dataset streaming#1003adithya-s-k wants to merge 26 commits into
adithya-s-k wants to merge 26 commits into
Conversation
A multimodal (vision + text) single-step (bandit) RL environment for image -> LaTeX transcription, backed by a Hugging Face dataset (default unsloth/LaTeX_OCR) and served through OpenEnv. Built on the newly introduced Task API (huggingface#726) and adds dataset streaming for datasets too large to materialize: - Task API over the dataset: list_splits / num_tasks / get_task / get_task_range. - Two access modes (LATEX_OCR_MODE): - materialize: split loaded + indexed; reset(split, index) random access. - stream: sequential cursor, NO full download; observations carry progress (index / total / remaining / pct_done); num_tasks from dataset metadata only; no-repeat within a pass. For TB-scale datasets. - Multimodal: image input + LaTeX output, gradeable with any vision-LLM policy. - Reward (LatexOCRRubric): 0.8*(1-CER) + 0.2*exact_match vs the hidden ground truth, whitespace-insensitive; pure-Python edit distance (no deps). - Custom Gradio "Try it" tab; typed client + Task API helpers; validate.py driver. - Unit tests for the rubric (tests/envs/test_latex_ocr_rubric.py).
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Pull request overview
This PR adds a new latex_ocr_env OpenEnv environment: a dataset-backed, single-step (bandit) image → LaTeX transcription task that exposes the Task API and supports dataset streaming for large datasets, along with a small rubric-focused test suite and an end-to-end validation script.
Changes:
- Introduces
envs/latex_ocr_env/(server, client, models, Dockerfile, OpenEnv manifest, docs, validate driver). - Implements a pure-Python LaTeX OCR reward rubric (CER + exact-match bonus) and adds unit tests for the rubric.
- Adds streaming mode support (sequential cursor + progress fields) in the environment implementation.
Alignment Review Report
Automated Checks
- Lint: NOT RUN — unable to execute repo hook scripts in this review environment.
- Debug code: NOT RUN (hook) — manual review did not find breakpoints/pdb; prints are present in
validate.py(expected for a driver script) and UI messaging.
Open RFCs Context
- RFC 002 (In Review): Task provider methods are intended for metadata/discovery and should be side-effect-free; Task API method semantics (e.g.,
list_tasksreturns all tasks). - RFC 004 (Implemented in code; design doc present): Rubric system exists in core; environments can expose
env.rubricfor introspection.
Tier 1: Fixes Required
-
envs/latex_ocr_env/server/latex_ocr_environment.py— streaming exhaustion can allow stale-target scoring (needs_target=None/_done=True). -
envs/latex_ocr_env/server/latex_ocr_environment.py— stream cursor/index bookkeeping is off-by-one and inconsistent across reset/step/task IDs. -
envs/latex_ocr_env/server/latex_ocr_environment.py— stream modereset(index=...)is silently ignored (should raise). -
envs/latex_ocr_env/server/latex_ocr_environment.py— stream modelist_tasks()returns a truncated preview, which conflicts with Task API expectations. -
tests/envs/test_latex_ocr_rubric.py— add a defensive assertion around dynamic import spec/loader to improve failure clarity.
Tier 2: Alignment Discussion
ALIGNMENT FLAG: list_tasks() semantics in stream mode (preview vs “all tasks”)
- Principle/RFC at stake: RFC 002 (TaskProvider semantics: discovery APIs;
list_tasks= all specs) - The concern: Returning a truncated list can mislead downstream tooling into thinking the split only has 100 tasks; raising
NotImplementedError(or otherwise making truncation explicit) is safer. - Suggested reviewer:
@darktex
ALIGNMENT FLAG: Custom rubric class vs core Rubric system introspection
- Principle/RFC at stake: RFC 004 (rubrics live inside environments; introspection via
env.rubric) - The concern:
LatexOCRRubricis implemented as a standalone helper rather than a coreRubric(openenv.core.rubrics.base.Rubric), which may reduce standard reward introspection/observability. - Suggested reviewer:
@darktex
Reviewed changes
Copilot reviewed 14 out of 15 changed files in this pull request and generated 7 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/envs/test_latex_ocr_rubric.py | Adds unit tests for the OCR reward rubric (pure-Python import path). |
| envs/latex_ocr_env/init.py | Exposes the typed client and models for the new env package. |
| envs/latex_ocr_env/client.py | Adds a typed EnvClient wrapper plus Task API HTTP helpers. |
| envs/latex_ocr_env/models.py | Defines Action/Observation wire models (including streaming progress fields). |
| envs/latex_ocr_env/openenv.yaml | Declares Space/runtime metadata and config defaults for deployment. |
| envs/latex_ocr_env/pyproject.toml | Adds per-env packaging and dependency declarations (incl. optional VLM driver deps). |
| envs/latex_ocr_env/README.md | Documents usage, Task API endpoints, and configuration env vars. |
| envs/latex_ocr_env/validate.py | Adds an end-to-end driver that can optionally call a VLM via the HF router. |
| envs/latex_ocr_env/server/init.py | Declares the server subpackage. |
| envs/latex_ocr_env/server/app.py | Creates the FastAPI app via core create_app, including a custom Gradio tab. |
| envs/latex_ocr_env/server/latex_ocr_environment.py | Implements the environment logic, Task API methods, and stream/materialize modes. |
| envs/latex_ocr_env/server/rubric.py | Implements the LaTeX OCR reward computation (CER + exact-match bonus). |
| envs/latex_ocr_env/server/gradio_ui.py | Adds a custom “Try it” Gradio playground that uses the live environment instance. |
| envs/latex_ocr_env/server/Dockerfile | Adds a container build for the environment server runtime. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+147
to
+153
| if self.mode == "stream": | ||
| # Positional stubs only; do not enumerate huge splits. | ||
| preview = min(n if n > 0 else 0, 100) | ||
| return [ | ||
| {"id": f"{split}-{i}", "index": i, "sequential": True} | ||
| for i in range(preview) | ||
| ] |
Comment on lines
+190
to
+194
| self._done = False | ||
| self._state = State(episode_id=episode_id or str(uuid4()), step_count=0) | ||
| if self.mode == "stream": | ||
| return self._reset_stream(split) | ||
| return self._reset_materialize(split, index, seed) |
Comment on lines
+242
to
+245
| except StopIteration: | ||
| self._exhausted = True | ||
| return LatexOCRObservation( | ||
| done=True, |
Comment on lines
+254
to
+257
| self._cursor += 1 | ||
| self._target = str(row[TEXT_COLUMN]) | ||
| self._current_split, self._current_index = split, self._cursor | ||
| remaining = (total - self._cursor) if total > 0 else -1 |
Comment on lines
+263
to
+266
| split=split, | ||
| index=self._cursor, | ||
| task_id=f"{split}-stream-{self._cursor}", | ||
| total=total, |
Comment on lines
+28
to
+32
| _spec = importlib.util.spec_from_file_location("latex_ocr_rubric", _RUBRIC_PATH) | ||
| rubric = importlib.util.module_from_spec(_spec) | ||
| # Register before exec so dataclasses in the module can resolve their module. | ||
| sys.modules[_spec.name] = rubric | ||
| _spec.loader.exec_module(rubric) |
Comment on lines
+237
to
+240
| def _reset_stream(self, split: str) -> LatexOCRObservation: | ||
| self._ensure_stream(split) | ||
| total = _split_total(self.dataset_name, split) | ||
| try: |
… 16) Pass max_concurrent_envs to create_app so multiple rollouts / UI users can run concurrent sessions; overridable via the LATEX_OCR_MAX_SESSIONS env var.
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 14 out of 15 changed files in this pull request and generated 2 comments.
Comments suppressed due to low confidence (3)
envs/latex_ocr_env/server/latex_ocr_environment.py:245
- In stream mode, when the dataset iterator is exhausted, reset() returns a done observation but does not mark the episode as terminated or clear the previous target. A subsequent step() call can incorrectly grade against the previous task instead of erroring out.
except StopIteration:
self._exhausted = True
return LatexOCRObservation(
done=True,
envs/latex_ocr_env/server/latex_ocr_environment.py:257
- Streaming cursor bookkeeping is off-by-one: _cursor is incremented before setting observation.index/_current_index, so the first streamed sample reports index=1. This also makes task_id and progress inconsistent with materialize mode (0-based row indices).
self._cursor += 1
self._target = str(row[TEXT_COLUMN])
self._current_split, self._current_index = split, self._cursor
remaining = (total - self._cursor) if total > 0 else -1
envs/latex_ocr_env/server/latex_ocr_environment.py:292
- In stream mode, reset() sets task_id to "{split}-stream-{index}" but step() currently reports "{split}-{index}", which can collide with materialize-mode IDs and makes it hard to correlate reset/step logs for the same task.
task_id=f"{self._current_split}-{self._current_index}",
Comment on lines
+169
to
+171
| n = self.num_tasks(split) | ||
| start = 0 if start is None else start | ||
| stop = n if stop is None else (min(stop, n) if n > 0 else stop) |
| [`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`, | ||
| where `CER` is the normalized character edit distance over whitespace-collapsed |
…ponents
Refactor LatexOCRRubric into a weighted sum of named components (each in [0,1],
weights renormalized -> reward in [0,1], tunable via env vars):
- edit_similarity (1 - CER) LATEX_OCR_W_EDIT (0.6)
- exact_match LATEX_OCR_W_EXACT (0.2)
- structural_validity (balanced delimiters + parseable LATEX_OCR_W_STRUCT (0.1)
LaTeX via optional pylatexenc)
- length_format (length agreement + no code fences) LATEX_OCR_W_LENFMT (0.1)
Per-component scores are surfaced in observation.reward_components and metadata
for training introspection. Empty prediction vs non-empty target scores 0 (no
floor credit). pylatexenc is optional (structural check degrades to a balance
check without it). Tests + docs updated.
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 14 out of 15 changed files in this pull request and generated 3 comments.
Comments suppressed due to low confidence (3)
envs/latex_ocr_env/server/latex_ocr_environment.py:270
- In stream mode, _reset_stream() increments _cursor before assigning index/task_id/current_index, so the first emitted task is numbered 1 (while Task API indices are 0-based). This also makes remaining/pct_done bookkeeping harder to reason about. Suggest keeping observation.index and _current_index 0-based and deriving pct_done from “consumed” = index+1.
self._cursor += 1
self._target = str(row[TEXT_COLUMN])
self._current_split, self._current_index = split, self._cursor
remaining = (total - self._cursor) if total > 0 else -1
return LatexOCRObservation(
done=False,
image_base64=_encode_image(row[IMAGE_COLUMN]),
image_format="png",
prompt=DEFAULT_PROMPT,
split=split,
index=self._cursor,
task_id=f"{split}-stream-{self._cursor}",
total=total,
remaining=remaining,
pct_done=round(self._cursor / total, 6) if total > 0 else 0.0,
exhausted=False,
)
envs/latex_ocr_env/server/latex_ocr_environment.py:293
- In stream mode, reset() returns task_id like "{split}-stream-{i}", but step() returns task_id "{split}-{i}". That makes task_id unstable across the episode and can confuse clients/logging. Suggest preserving the stream prefix in step() when mode=="stream".
reward=result.reward,
split=self._current_split,
index=self._current_index,
task_id=f"{self._current_split}-{self._current_index}",
predicted_latex=prediction,
tests/envs/test_latex_ocr_rubric.py:32
- The dynamic import in this test assumes spec_from_file_location() always returns a spec with a loader. If it fails (e.g., path issues), the error will be a less-informative AttributeError. Adding an explicit check will fail fast with a clearer message.
_spec = importlib.util.spec_from_file_location("latex_ocr_rubric", _RUBRIC_PATH)
rubric = importlib.util.module_from_spec(_spec)
# Register before exec so dataclasses in the module can resolve their module.
sys.modules[_spec.name] = rubric
_spec.loader.exec_module(rubric)
Comment on lines
+92
to
+96
| builder = load_dataset_builder( | ||
| dataset_name, name=CONFIG, token=os.environ.get("HF_TOKEN") | ||
| ) | ||
| info = builder.info.splits.get(split) | ||
| return int(info.num_examples) if info and info.num_examples else -1 |
Comment on lines
+72
to
+75
| gr.Markdown( | ||
| "Transcribe the image into **LaTeX**, then score it. Reward = " | ||
| "`0.8·(1−CER) + 0.2·exact_match` against the hidden ground truth." | ||
| ) |
| ``reset()`` pulls the *next* sample; every observation carries progress | ||
| (``index``, ``total``, ``remaining``, ``pct_done``). No-repeat within a pass; | ||
| ``reset(index=i)`` random access is unsupported by design. Best for very large | ||
| / TB-scale datasets. See STREAMING.md. |
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 14 out of 15 changed files in this pull request and generated 1 comment.
Comments suppressed due to low confidence (3)
envs/latex_ocr_env/server/latex_ocr_environment.py:194
- In stream mode, reset() currently ignores the provided
indexargument. That makesreset(split, index=...)appear supported but behave differently than requested (and contradicts the module docstring that says random access is unsupported). Consider rejectingindexexplicitly in stream mode so callers don’t get silent surprises.
if split not in self.list_splits():
split = self.list_splits()[0]
self._done = False
self._state = State(episode_id=episode_id or str(uuid4()), step_count=0)
if self.mode == "stream":
return self._reset_stream(split)
return self._reset_materialize(split, index, seed)
envs/latex_ocr_env/server/latex_ocr_environment.py:258
- Stream mode uses a 1-based
index/cursor and emits task IDs that don’t match the Task API’sget_task()IDs (and changes the task_id format between reset and step). This makes tasks hard to correlate across Task API vs episode results and introduces an off-by-one in progress fields.
self._cursor += 1
self._target = str(row[TEXT_COLUMN])
self._current_split, self._current_index = split, self._cursor
remaining = (total - self._cursor) if total > 0 else -1
return LatexOCRObservation(
envs/latex_ocr_env/server/gradio_ui.py:75
- The UI text hard-codes a reward formula (
0.8·(1−CER) + 0.2·exact_match) that doesn’t match the rubric’s actual default weights (0.6/0.2/0.1/0.1). This can confuse users when they see rewards/components that don’t align with the stated formula.
gr.Markdown(
"Transcribe the image into **LaTeX**, then score it. Reward = "
"`0.8·(1−CER) + 0.2·exact_match` against the hidden ground truth."
)
Comment on lines
+92
to
+96
| rewards = [] | ||
| seen_targets = [] | ||
| for i in range(min(args.num, n)): | ||
| # In stream mode `index` is ignored — reset() pulls the next sample. | ||
| result = env.reset(split=args.split, index=i) |
The composable structural-validity + length/format components introduced
discontinuous (partly binary) terms that made the reward spiky and training
unstable. Roll the rubric back to the smooth dense signal:
reward = 0.8 * (1 - CER) + 0.2 * exact_match
Removes the extra components, the reward_components field, weight env vars, and
the pylatexenc dependency. Keeps the Gradio UI fix that reads reward from the
top-level step result (observation.reward is None after server serialization).
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 14 out of 15 changed files in this pull request and generated no new comments.
Comments suppressed due to low confidence (3)
envs/latex_ocr_env/server/latex_ocr_environment.py:246
- In streaming mode, when the iterator is exhausted, the environment returns a terminal observation but leaves
_target/_doneunchanged. If a caller mistakenly callsstep()after an exhaustedreset(), it may score against a stale target from a previous task. Clear_targetand mark the episode done in the StopIteration path (and optionally update current split/index) to prevent reuse of stale state.
except StopIteration:
self._exhausted = True
return LatexOCRObservation(
done=True,
split=split,
envs/latex_ocr_env/server/latex_ocr_environment.py:293
reset()(stream mode) returns task IDs like"{split}-stream-{n}", butstep()currently reportstask_idas"{split}-{n}". This makes task identifiers inconsistent within a single episode and can collide with materialize-mode IDs. Use the stream-prefixed task_id format whenmode == "stream".
split=self._current_split,
index=self._current_index,
task_id=f"{self._current_split}-{self._current_index}",
predicted_latex=prediction,
tests/envs/test_latex_ocr_rubric.py:32
- This test dynamically loads
rubric.pyviaspec_from_file_location, but doesn't guard against_spec(or_spec.loader) beingNone. If spec creation fails, the test will error with an AttributeError rather than a clear assertion failure. Add an explicit assertion so failures are easier to diagnose.
_spec = importlib.util.spec_from_file_location("latex_ocr_rubric", _RUBRIC_PATH)
rubric = importlib.util.module_from_spec(_spec)
# Register before exec so dataclasses in the module can resolve their module.
sys.modules[_spec.name] = rubric
_spec.loader.exec_module(rubric)
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 14 out of 15 changed files in this pull request and generated no new comments.
Comments suppressed due to low confidence (3)
envs/latex_ocr_env/server/latex_ocr_environment.py:258
- In stream mode,
_cursoris incremented before being assigned toindex/task_id/_current_index, making indices 1-based and causingtask_idonreset()(e.g.split-stream-1) to disagree with the Task API’sget_task(split, 0)(split-0) and withstep()’s returnedtask_id(split-<index>). This breaks stable task identification and makes progress reporting off-by-one.
self._cursor += 1
self._target = str(row[TEXT_COLUMN])
self._current_split, self._current_index = split, self._cursor
remaining = (total - self._cursor) if total > 0 else -1
return LatexOCRObservation(
envs/latex_ocr_env/server/latex_ocr_environment.py:172
get_task_range()behaves incorrectly whennum_tasks()returns an unknown size (e.g. stream mode returns-1): ifstopis omitted,stopbecomes-1, sorange(start, stop)is empty and callers can’t page tasks at all. This should either require an explicitstopwhenn < 0, or choose a sensible default page size.
n = self.num_tasks(split)
start = 0 if start is None else start
stop = n if stop is None else (min(stop, n) if n > 0 else stop)
return [
tests/envs/test_latex_ocr_rubric.py:32
- The test dynamically loads
rubric.pyviaspec_from_file_location(), but doesn’t guard against_specor_spec.loaderbeingNone(which can happen if the path is wrong or import machinery can’t create a loader). This would raise an AttributeError later and obscure the real failure.
_spec = importlib.util.spec_from_file_location("latex_ocr_rubric", _RUBRIC_PATH)
rubric = importlib.util.module_from_spec(_spec)
# Register before exec so dataclasses in the module can resolve their module.
sys.modules[_spec.name] = rubric
_spec.loader.exec_module(rubric)
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 18 out of 19 changed files in this pull request and generated 2 comments.
Comments suppressed due to low confidence (4)
envs/latex_ocr_env/server/latex_ocr_environment.py:323
- In stream mode, the cursor is incremented before being recorded in
index/task_id, making the first task report index=1 (and the last task index==total). This is inconsistent with the Task API stubs (get_task/list_tasksare 0-based) and makes progress/accounting off by one.
self._cursor += 1
self._target = str(row[TEXT_COLUMN])
self._current_split, self._current_index = split, self._cursor
self._current_task_id = f"{split}-stream-{self._cursor}"
remaining = (total - self._cursor) if total > 0 else -1
envs/latex_ocr_env/server/latex_ocr_environment.py:352
- In stream mode,
step()uses_split_total()(metadata) fortotal, butreset()usesnum_tasks()(which also honorsLATEX_OCR_MAX_ROWS). WhenLATEX_OCR_MAX_ROWSis set, this makestotal/remaining/pct_doneinconsistent between reset and step.
total = (
_split_total(self.dataset_name, self._current_split)
if self.mode == "stream"
else -1
)
envs/latex_ocr_env/server/gradio_ui.py:75
- The Gradio UI text hardcodes reward weights (0.8/0.2), but the environment defaults to
exact_weight=0.4(edit=0.6, exact=0.4) and is env-var configurable. This is user-facing documentation drift inside the UI.
gr.Markdown(
"Transcribe the image into **LaTeX**, then score it. Reward = "
"`0.8·(1−CER) + 0.2·exact_match` against the hidden ground truth."
)
examples/grpo_latex_ocr/README.md:39
- This README claims the reward includes “structural validity”, but the current rubric implementation is edit similarity + exact match (plus the length guard). This mismatch will confuse readers and makes the example documentation inaccurate.
See [`envs/latex_ocr_env`](../../envs/latex_ocr_env) for the full environment: a dataset-backed,
single-step (bandit) RL task with a weighted, server-side reward (edit similarity, exact match,
structural validity, length/format). It ships `train` and `test` splits.
Comment on lines
+211
to
+214
| return [ | ||
| {"id": f"{split}-{i}", "index": i, "split": split} | ||
| for i in range(start, stop) | ||
| ] |
Comment on lines
+133
to
+164
| def grade(self, prediction: str, target: str) -> GradeResult: | ||
| raw = prediction or "" | ||
| # Clean the raw completion (fences, $-delimiters) before scoring; LaTeX spacing is | ||
| # cosmetic, so score on the whitespace-stripped form: an exact match => CER 0. | ||
| cleaned = clean_latex(raw) | ||
| pred_canon = _strip_all_whitespace(normalize_latex(cleaned)) | ||
| target_norm = normalize_latex(target) | ||
| target_canon = _strip_all_whitespace(target_norm) | ||
|
|
||
| exact = pred_canon == target_canon | ||
| if not target_canon: | ||
| cer = 0.0 if not pred_canon else 1.0 | ||
| else: | ||
| distance = levenshtein(pred_canon, target_canon) | ||
| cer = min(1.0, distance / max(len(pred_canon), len(target_canon))) | ||
|
|
||
| similarity = 1.0 - cer | ||
| reward = (1.0 - self.exact_weight) * similarity + self.exact_weight * float( | ||
| exact | ||
| ) | ||
|
|
||
| # Length guard: penalize predictions whose RAW length (measured before whitespace is | ||
| # stripped, so padding counts) far exceeds the target's. Reward decays linearly with | ||
| # the excess and reaches 0 at twice the allowance. | ||
| length_factor = 1.0 | ||
| if self.overlong_ratio > 0: | ||
| allowed = max(self.overlong_floor, self.overlong_ratio * len(target_norm)) | ||
| if len(raw) > allowed: | ||
| over = (len(raw) - allowed) / allowed | ||
| length_factor = max(0.0, 1.0 - over) | ||
| reward *= length_factor | ||
|
|
Add held-out eval results + train/eval curves for four VLMs trained against the latex_ocr env (Qwen3-VL-2B, Qwen3.5-2B, GLM-OCR, Gemma-4-E2B), showing stable, monotonic improvement (deltas +3% to +48%).
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 18 out of 20 changed files in this pull request and generated no new comments.
Comments suppressed due to low confidence (6)
envs/latex_ocr_env/server/latex_ocr_environment.py:352
- In stream mode,
step()reportstotal/remaining/pct_doneusing_split_total()(metadata) and ignoresLATEX_OCR_MAX_ROWS, butreset()usesnum_tasks()which does honor the cap. This can make progress jump or never reach 100% when a cap is set.
total = (
_split_total(self.dataset_name, self._current_split)
if self.mode == "stream"
else -1
)
envs/latex_ocr_env/server/latex_ocr_environment.py:214
get_task_range()omits thesequentialflag in stream mode, butlist_tasks()/get_task()include it. Returning consistent task metadata helps clients understand random access is unsupported in stream mode.
return [
{"id": f"{split}-{i}", "index": i, "split": split}
for i in range(start, stop)
]
envs/latex_ocr_env/server/gradio_ui.py:75
- The UI text hard-codes a
0.8/0.2reward split, but the environment defaults toexact_weight=0.4(and it’s env-var tunable). This is user-facing documentation and will mislead users about reward semantics.
gr.Markdown(
"Transcribe the image into **LaTeX**, then score it. Reward = "
"`0.8·(1−CER) + 0.2·exact_match` against the hidden ground truth."
)
examples/grpo_latex_ocr/README.md:3
- The Colab badge links to a personal fork/branch (
adithya-s-k/OpenEnv+add-latex-ocr-...). Once merged, this URL is likely to 404 and the example becomes unusable from main.
[](https://colab.research.google.com/github/adithya-s-k/OpenEnv/blob/add-latex-ocr-multimodal-streaming-env/examples/grpo_latex_ocr/grpo_latex_ocr_tutorial.ipynb)
examples/grpo_latex_ocr/README.md:39
- This README claims the environment reward includes “structural validity”, but the rubric implemented in this PR is CER + exact-match with a length guard (no structural-validity term). Keeping this accurate matters for users interpreting the results table.
See [`envs/latex_ocr_env`](../../envs/latex_ocr_env) for the full environment: a dataset-backed,
single-step (bandit) RL task with a weighted, server-side reward (edit similarity, exact match,
structural validity, length/format). It ships `train` and `test` splits.
examples/grpo_latex_ocr/grpo_latex_ocr_tutorial.ipynb:41
- The notebook installs the environment from a personal fork/branch. After merge, this should reference the canonical repo (or PyPI) so the tutorial works for readers.
" \"openenv-latex_ocr_env @ git+https://github.com/adithya-s-k/OpenEnv.git@add-latex-ocr-multimodal-streaming-env#subdirectory=envs/latex_ocr_env\" \\\n",
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 18 out of 20 changed files in this pull request and generated 1 comment.
Comments suppressed due to low confidence (7)
envs/latex_ocr_env/server/latex_ocr_environment.py:323
- In stream mode,
_cursoris incremented before it is used as the task index, which makes the first streamed task reportindex=1(andtask_id=...-1) even though Task API indices are 0-based. This also makesremaining/pct_doneoff by one and inconsistent withget_task_range/list_tasksstubs.
self._cursor += 1
self._target = str(row[TEXT_COLUMN])
self._current_split, self._current_index = split, self._cursor
self._current_task_id = f"{split}-stream-{self._cursor}"
remaining = (total - self._cursor) if total > 0 else -1
envs/latex_ocr_env/server/latex_ocr_environment.py:352
- In stream mode,
step()reportstotalusing_split_total()(full dataset metadata), butreset()reportstotalusingnum_tasks()(which honorsLATEX_OCR_MAX_ROWS). With a dev cap set, the denominator can change mid-episode, andremaining/pct_donebecome inconsistent.
total = (
_split_total(self.dataset_name, self._current_split)
if self.mode == "stream"
else -1
)
envs/latex_ocr_env/server/latex_ocr_environment.py:214
get_task_range()doesn't include thesequentialflag in stream mode, even thoughget_task()andlist_tasks()do. This makes Task API responses inconsistent across endpoints for stream mode clients.
This issue also appears in the following locations of the same file:
- line 319
- line 348
def get_task_range(
self, split: str, start: int | None = None, stop: int | None = None
) -> list[dict[str, Any]]:
n = self.num_tasks(split)
start = 0 if start is None else start
stop = n if stop is None else (min(stop, n) if n > 0 else stop)
# In stream mode `n` comes from metadata and can be enormous; cap the number
# of generated stubs so an unbounded range can't OOM the process.
if self.mode == "stream" and stop - start > _STREAM_RANGE_CAP:
logger.warning(
"get_task_range: capping stream range %d..%d to %d stubs",
start,
stop,
_STREAM_RANGE_CAP,
)
stop = start + _STREAM_RANGE_CAP
return [
{"id": f"{split}-{i}", "index": i, "split": split}
for i in range(start, stop)
]
envs/latex_ocr_env/server/gradio_ui.py:75
- The Gradio UI text hard-codes a reward formula with 0.8/0.2 weights, but the environment default is
exact_weight=0.4(i.e., 0.6/0.4). This is misleading for users trying to interpret scores.
gr.Markdown(
"Transcribe the image into **LaTeX**, then score it. Reward = "
"`0.8·(1−CER) + 0.2·exact_match` against the hidden ground truth."
)
examples/grpo_latex_ocr/README.md:39
- This README claims the env reward includes “structural validity”, but the rubric in this PR is edit similarity + exact match (+ length guard). The extra criteria aren’t implemented, so the example description is currently inaccurate.
See [`envs/latex_ocr_env`](../../envs/latex_ocr_env) for the full environment: a dataset-backed,
single-step (bandit) RL task with a weighted, server-side reward (edit similarity, exact match,
structural validity, length/format). It ships `train` and `test` splits.
examples/grpo_latex_ocr/README.md:3
- The Colab badge URL points to a fork/feature branch (
adithya-s-k/OpenEnv/blob/add-latex-ocr-multimodal-streaming-env/...). After merge, this should link to the notebook in the mainhuggingface/OpenEnvrepo (or a stable tag) so the badge stays valid.
[](https://colab.research.google.com/github/adithya-s-k/OpenEnv/blob/add-latex-ocr-multimodal-streaming-env/examples/grpo_latex_ocr/grpo_latex_ocr_tutorial.ipynb)
A beginner-friendly tutorial that uses **GRPO (Group Relative Policy Optimization)** to teach a
**vision-language model** to transcribe images of math formulas into **LaTeX** — with the images
*and* the reward supplied by the [`latex_ocr_env`](../../envs/latex_ocr_env) OpenEnv environment.
examples/grpo_latex_ocr/grpo_latex_ocr_tutorial.ipynb:41
- The notebook installs the environment from a personal fork/branch (
adithya-s-k/OpenEnv@add-latex-ocr-multimodal-streaming-env). After this PR merges, that reference will drift or disappear; examples in-repo should point at the canonicalhuggingface/OpenEnv(or a tagged release / commit SHA).
"!pip install -q \"trl @ git+https://github.com/huggingface/trl.git\" \\\n",
" \"openenv-latex_ocr_env @ git+https://github.com/adithya-s-k/OpenEnv.git@add-latex-ocr-multimodal-streaming-env#subdirectory=envs/latex_ocr_env\" \\\n",
" peft trackio ipywidgets fastmcp websockets jmespath hf_transfer"
Comment on lines
+143
to
+148
| if not target_canon: | ||
| cer = 0.0 if not pred_canon else 1.0 | ||
| else: | ||
| distance = levenshtein(pred_canon, target_canon) | ||
| cer = min(1.0, distance / max(len(pred_canon), len(target_canon))) | ||
|
|
12 tasks
Member
|
Really nice work, and a great first consumer of the Task API. I read through the full diff and the earlier bot review rounds: the stream-exhaustion, MAX_ROWS, OOM-cap, PNG-normalization and unknown-split fixes all landed cleanly, and the whitespace-padding length guard is well designed and well tested. The points below are new, none of them were raised in the previous rounds. 1. Stream mode duplicates data across concurrent sessions (design, worth documenting at minimum)With For the TB-scale use case this env is pitched at, that's the main practical caveat. Suggestions, in increasing order of effort:
2. Gradio copy drifted from the new default reward weights
Same drift in 3. Client
|
Bounded edit-distance grading and a corrected Task API: Python slice semantics in get_task_range, split validation on every Task API method, and list_tasks refusing to pretend it can enumerate a streamed split.
…enominator mid-episode Two stream-mode reporting bugs, both raised in review. num_tasks applied LATEX_OCR_MAX_ROWS only when dataset metadata reported a positive total. A dataset whose split has no num_examples returns -1, and the cap was then ignored, so callers were told the size was unknown even though the cursor stops at exactly the cap. The cap is the answer in that case. step recomputed total from _split_total, which does not know about the cap, while reset used num_tasks, which does. With a cap configured a trainer saw 1/50 from reset and 1/900 from step for the same split in the same episode.
The client built its StepResult with info=, but StepResult exposes metadata. Every step therefore raised TypeError inside the try, fell through to the fallback, and dropped whatever the server sent. metadata is now passed directly, with info read as a fallback for the older server shape. create_app only mounts the Gradio UI when ENABLE_WEB_INTERFACE is set, and nothing set it: not the Dockerfile, not app.py. The documented Try it tab was dark on Docker, on a Space and on the local run path in app.py's own docstring. Set via setdefault, so ENABLE_WEB_INTERFACE=false still wins.
Its two parents were already merged here individually, so it carries no content of its own; -s ours keeps this tree exactly as tested.
Generated by scripts/sync_env_docs.py --fix. The stub mirrors the env README, and Ben's Task API note in the README left it stale, which is what check-env-docs fails on.
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.
There are 2 total unresolved issues (including 1 from previous review).
❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.
Reviewed by Cursor Bugbot for commit 9fc674b. Configure here.
| prompt=DEFAULT_PROMPT, | ||
| split=split, | ||
| index=self._cursor, | ||
| task_id=f"{split}-stream-{self._cursor}", |
There was a problem hiding this comment.
Stream task IDs mismatch Task API
Medium Severity
In stream mode, episode observations use 1-based index values and task_ids like train-stream-1, while get_task / get_task_range emit 0-based specs with ids like train-0. Trainers cannot join Task API discovery metadata to episode results for the same sequential position.
Additional Locations (1)
Reviewed by Cursor Bugbot for commit 9fc674b. Configure here.
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 20 out of 22 changed files in this pull request and generated no new comments.
Suppressed comments (4)
envs/latex_ocr_env/server/rubric.py:153
- The length guard can raise ZeroDivisionError when
allowedbecomes 0 (e.g. empty target plusLATEX_OCR_OVERLONG_FLOOR=0withoverlong_ratio>0). This would crash grading for that configuration; guard the division by checkingallowed > 0(or validateoverlong_floor >= 1).
if self.overlong_ratio > 0:
allowed = max(self.overlong_floor, self.overlong_ratio * len(target_norm))
if len(raw) > allowed:
over = (len(raw) - allowed) / allowed
length_factor = max(0.0, 1.0 - over)
envs/latex_ocr_env/server/gradio_ui.py:75
- The UI text hard-codes the reward weights as
0.8/0.2, but the rubric/environment default isexact_weight=0.4(and it’s configurable via env var). This makes the “Try it” tab misleading; it should describe the formula in terms ofexact_weight(or mention the default 0.4).
gr.Markdown(
"Transcribe the image into **LaTeX**, then score it. Reward = "
"`0.8·(1−CER) + 0.2·exact_match` against the hidden ground truth."
)
examples/grpo_latex_ocr/README.md:3
- The Colab badge URL points to a personal fork/feature branch (
adithya-s-k/OpenEnv/blob/add-latex-ocr-...). After merge this will be stale/broken for users reading the repo docs; it should point at the canonicalhuggingface/OpenEnvdefault branch (or use a relative link pattern).
[](https://colab.research.google.com/github/adithya-s-k/OpenEnv/blob/add-latex-ocr-multimodal-streaming-env/examples/grpo_latex_ocr/grpo_latex_ocr_tutorial.ipynb)
examples/grpo_latex_ocr/grpo_latex_ocr_tutorial.ipynb:41
- This notebook installs the env package from a personal fork/branch (
adithya-s-k/OpenEnv@add-latex-ocr-...). That makes the tutorial non-reproducible once the branch disappears and also bypasses the canonical upstream; it should install fromhuggingface/OpenEnv(ideally pinned to a tag/commit).
" \"openenv-latex_ocr_env @ git+https://github.com/adithya-s-k/OpenEnv.git@add-latex-ocr-multimodal-streaming-env#subdirectory=envs/latex_ocr_env\" \\\n",
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 20 out of 22 changed files in this pull request and generated no new comments.
Suppressed comments (2)
envs/latex_ocr_env/server/gradio_ui.py:75
- The UI text hard-codes
0.8·(1−CER) + 0.2·exact_match, but the environment’s default reward weighting isexact_weight=0.4(so edit is 0.6 / exact is 0.4). This is also inconsistent with the formula documented elsewhere in this PR; consider expressing the formula in terms ofexact_weightinstead of fixed constants.
gr.Markdown(
"Transcribe the image into **LaTeX**, then score it. Reward = "
"`0.8·(1−CER) + 0.2·exact_match` against the hidden ground truth."
)
envs/latex_ocr_env/validate.py:106
validate.pyswitches to stream-mode on any exception fromreset(split, index=...), which can hide real failures (e.g., connection errors, bad split). Consider only falling back when the error clearly indicates indexed reset is unsupported.
try:
result = env.reset(split=args.split, index=i)
except Exception:
# stream-mode server rejects random index -> pull sequentially
stream_mode = True
result = env.reset(split=args.split)
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Closes #1002
What
Adds
envs/latex_ocr_env/— a multimodal (vision + text) single-step(bandit) RL environment for image → LaTeX transcription, built on the
newly introduced Task API (#726) and adding dataset streaming for
datasets too large to materialize. Backed by a Hugging Face dataset (default
unsloth/LaTeX_OCR).Why it matters
of
list_splits/num_tasks/get_task/get_task_range.TB-scale datasets (millions of rows) can be served from a small Space.
Highlights
list_splits,num_tasks,get_task,get_task_range.LATEX_OCR_MODE):materialize(default) — split loaded + indexed;reset(split, index)random access. Best when the dataset fits on disk.
stream— sequential cursor over a streamed split (no full download);observations carry
index / total / remaining / pct_done;num_tasksfromdataset metadata only; no-repeat within a pass. For TB-scale datasets.
LatexOCRRubric) —0.8·(1−CER) + 0.2·exact_matchagainst thehidden ground truth, whitespace-insensitive; pure-Python edit distance.
(via the
gradio_builderhook).reset/step+ Task API helpers;validate.pyend-to-end driver (optionally drives a vision-LLM policy via the HF Router).
Files
Testing
PYTHONPATH=src:envs uv run pytest tests/envs/test_latex_ocr_rubric.py— 8 passing.ruff format/ruff check/usort— clean.progress, no-repeat, and a real vision-LLM policy (mean reward ~0.77, incl.
exact matches).
Notes
openenv>=0.4.1.Environment+ new Task API.(image, latex)dataset works viaLATEX_OCR_DATASET+ column env vars.Note
Low Risk
Self-contained new env under
envs/with no core changes; main operational risk is HF dataset download/streaming and optionalHF_TOKENat runtime.Overview
Adds
latex_ocr_env, a new single-step (bandit) OpenEnv for image → LaTeX transcription on Hugging Face datasets (defaultunsloth/LaTeX_OCR), wired to the Task API (list_splits,num_tasks,get_task,get_task_range) plus a typedLatexOCREnvclient.Serving:
materializemode loads/indexes splits forreset(split, index);streammode walks a sequential cursor without a full download, with progress fields and caps on hugeget_task_rangestubs. Ground truth stays hidden untilstep; reward is server-sideLatexOCRRubric(CER + exact-match blend, whitespace-insensitive, optional length guard against padding hacks, fence/$stripping).Also ships FastAPI app (defaults
ENABLE_WEB_INTERFACEfor the Gradio Try it tab), Dockerfile, docs toctree entry,validate.py, and tests for rubric, Task API edge cases, and client metadata parsing.Reviewed by Cursor Bugbot for commit 9fc674b. Bugbot is set up for automated code reviews on this repo. Configure here.