Add LaTeX OCR: multimodal env on the new Task API + dataset streaming - #1003
Open
adithya-s-k wants to merge 19 commits into
Open
Add LaTeX OCR: multimodal env on the new Task API + dataset streaming#1003adithya-s-k wants to merge 19 commits into
adithya-s-k wants to merge 19 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.
| return [ | ||
| {"id": f"{split}-{i}", "index": i, "split": split} | ||
| for i in range(start, stop) | ||
| ] |
There was a problem hiding this comment.
Task range ignores slice semantics
Medium Severity
get_task_range builds range(start, stop) after only null-coalescing bounds, but the Task API contract requires Python slice semantics. Negative start/stop therefore emit invalid negative indexes (or an empty list) instead of counting from the end of the split, unlike the reference Task API implementations.
Reviewed by Cursor Bugbot for commit a9fa176. Configure here.
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 (6)
envs/latex_ocr_env/server/latex_ocr_environment.py:192
- In stream mode, if the dataset metadata lacks a split size (
num_tasksreturns -1) andstopis omitted,stopbecomes -1 andrange(start, stop)returns an empty list. Also, stream-mode task stubs returned here omit thesequentialflag thatlist_tasks()/get_task()include.
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.
envs/latex_ocr_env/server/latex_ocr_environment.py:313
- Stream-mode indexing is off-by-one:
_cursoris incremented before populatingindex/task_id, so the first streamed task reports index=1 and uses IDs that don't align with the Task API stubs (which start at 0). This also makesremaining/pct_donesemantics harder to reason about.
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:342
- In stream mode,
reset()reportstotalvianum_tasks()(which honorsLATEX_OCR_MAX_ROWS), butstep()reportstotalvia_split_total()(which ignores the cap). This makestotal/remaining/pct_doneinconsistent between reset and step whenLATEX_OCR_MAX_ROWSis set.
total = (
_split_total(self.dataset_name, self._current_split)
if self.mode == "stream"
else -1
)
examples/grpo_latex_ocr/README.md:3
- The Colab badge link points to a personal fork/feature branch, which will become stale once this PR merges. It should reference the canonical
huggingface/OpenEnvrepo (and typically themainbranch).
[](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 environment description lists reward components ("structural validity, length/format") that aren't implemented by
LatexOCRRubric(currently CER + exact match only). This can mislead readers about what the environment actually scores.
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 tutorial installs the environment package from a personal fork/feature branch, which will become stale and isn't the canonical source for this repository. It should point to
huggingface/OpenEnv(ideally a tag/commit) for long-term reproducibility.
" \"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",
Comment on lines
+61
to
+68
| total = o.get("total", -1) | ||
| if total and total > 0: | ||
| return ( | ||
| f"**Task `{o.get('task_id', '')}`** · position " | ||
| f"**{o.get('index')}/{total}** " | ||
| f"({o.get('pct_done', 0):.4%}) · remaining **{o.get('remaining')}**" | ||
| ) | ||
| return f"**Task `{o.get('task_id', '')}`** (index {o.get('index')})" |
Comment on lines
+108
to
+113
| prog = "" | ||
| if obs.total and obs.total > 0: | ||
| prog = ( | ||
| f" | progress: {obs.index}/{obs.total} " | ||
| f"({obs.pct_done:.4%}), remaining={obs.remaining}" | ||
| ) |
…server imports its UI)
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 no new comments.
Comments suppressed due to low confidence (5)
envs/latex_ocr_env/server/latex_ocr_environment.py:313
- In stream mode,
_cursoris incremented before being used forindex/task_id, which makes the first streamed task reportindex=1(andtask_idsuffix1) even though Task API indices are 0-based. This creates an off-by-one mismatch betweenget_task(split, 0)and the firstreset()result in stream mode.
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:342
- In stream mode,
step()uses_split_total(...)(raw metadata) fortotal, butreset()usesself.num_tasks()which appliesLATEX_OCR_MAX_ROWS. IfLATEX_OCR_MAX_ROWSis set,total/remaining/pct_donewill disagree between reset and step responses.
total = (
_split_total(self.dataset_name, self._current_split)
if self.mode == "stream"
else -1
)
examples/grpo_latex_ocr/README.md:3
- The Colab badge URL points at a personal fork/branch (
adithya-s-k/OpenEnv@add-latex-ocr-multimodal-streaming-env). After merge, that link will likely 404. It should point at the canonicalhuggingface/OpenEnvlocation for this notebook.
[](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, length/format", but the environment's rubric implementation is only edit similarity (CER) + exact match. This description should match the actual reward components to avoid misleading users.
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
openenv-latex_ocr_envfrom a personal fork/branch (adithya-s-k/OpenEnv@add-latex-ocr-multimodal-streaming-env). After merge, this will likely break. Install from the canonicalhuggingface/OpenEnvrepo (or a release tag) instead.
" \"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",
…ar tunable Default exact_weight 0.2 -> 0.4 so partial answers cap at 0.6 and exact matches are rewarded more strongly. Tunable at runtime via LATEX_OCR_EXACT_WEIGHT.
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 1 comment.
Comments suppressed due to low confidence (6)
envs/latex_ocr_env/server/latex_ocr_environment.py:197
- In stream mode, get_task_range() can return an empty list when num_tasks() is unknown (n <= 0). Specifically, when stop is omitted (None) and n is -1, stop becomes -1 and range(start, stop) yields no stubs, which breaks default Task API enumeration.
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(
envs/latex_ocr_env/server/latex_ocr_environment.py:345
- step() computes total via _split_total() in stream mode, which ignores LATEX_OCR_MAX_ROWS. This makes total/remaining/pct_done inconsistent between reset() (which uses num_tasks() and honors the cap) 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 UI hard-codes a
0.8/0.2reward split, but the environment default exact_weight is 0.4 (and is configurable). This text should match the actual rubric or describe the parameterized 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."
)
examples/grpo_latex_ocr/README.md:3
- This Colab badge links to a personal fork/feature branch. Other examples in this repo link to
github/huggingface/OpenEnv/blob/main/...; this one should too so it stays valid after merge.
[](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
- The reward description mentions "structural validity" and "length/format", but the LaTeX OCR environment rubric shown in this PR is edit-distance similarity + exact match only. This should be corrected to avoid misleading users.
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
openenv-latex_ocr_envfrom a personal fork/branch. After merge, this should point athuggingface/OpenEnv(like other notebooks) so the example remains reproducible (e.g.,git+https://github.com/huggingface/OpenEnv.git#subdirectory=envs/latex_ocr_env).
" \"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",
Comment on lines
+12
to
+16
| - **Reward**: `(1 - exact_weight) * (1 - CER) + exact_weight * exact_match`, | ||
| where `CER` is the normalized character edit distance over whitespace-stripped | ||
| LaTeX. Partial answers score in `[0, 0.8]`; 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. |
sergiopaniego
left a comment
sergiopaniego
left a comment
Member
There was a problem hiding this comment.
great one!
we need to solve some comments from automated PR reviews, many are legit.
some ideas:
- the CI is failing:
python scripts/sync_env_docs.py --fixshould fix it. - env hosting: the env could be moved the OpenEnv's org and in the notebook we could link it directly
- docs: we could add the env to
docs/source/environments.md
The whitespace-insensitive reward let a policy score 1.0 by emitting the correct answer followed by unlimited whitespace (the padding is stripped before scoring). In RL this shows up as completions drifting to the generation-length cap. Add a length_factor to the rubric that decays the reward once the RAW prediction exceeds overlong_ratio x the target length (default 4.0, floor 80 chars; tunable via LATEX_OCR_OVERLONG_RATIO / LATEX_OCR_OVERLONG_FLOOR, 0 disables). The rubric now also cleans raw completions (code fences, $-delimiters) itself and measures their true length, so it can score model output directly. Legit answers and normal LaTeX spacing are unaffected; padded completions go from 1.0 to ~0.
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",
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.
There are 5 total unresolved issues (including 4 from previous reviews).
❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.
Want higher recall? High effort reviews run extra passes and find more bugs. A team admin can switch effort levels in the Cursor dashboard.
Reviewed by Cursor Bugbot for commit 8ba41aa. Configure here.
| try: | ||
| return StepResult(**base, info=data.get("info", {})) | ||
| except TypeError: | ||
| return StepResult(**base) |
There was a problem hiding this comment.
Client drops StepResult metadata
Low Severity
_parse_result tries to construct StepResult with info=, but current core’s StepResult exposes metadata, not info. The TypeError fallback omits it, so top-level step metadata from the server is always dropped even though sibling env clients pass metadata= successfully.
Reviewed by Cursor Bugbot for commit 8ba41aa. Configure here.
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
|
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 environment under
envs/with no core changes; reward and streaming behavior are covered by unit tests, with main operational risk being HF dataset download size and memory in materialize mode.Overview
Introduces
envs/latex_ocr_env/, a new single-step (bandit) OpenEnv for image → LaTeX transcription backed by Hugging Face datasets (defaultunsloth/LaTeX_OCR).The server implements the Task API (
list_splits,num_tasks,get_task,get_task_range) and supportsmaterialize(indexed random access) vsstream(sequential cursor, metadata-only counts, progress fields on observations). Ground truth stays hidden onreset;LatexOCRRubricgrades server-side with CER + exact-match weighting, fence/$stripping, and an optional length guard against whitespace-padding hacks.Also adds a typed
LatexOCREnvclient (WebSocketreset/step+ HTTP Task helpers), FastAPI app with Gradio Try it tab, Docker packaging, docs/toc entry,validate.py, rubric unit tests, and a GRPO tutorial example that uses the env as the reward source.Reviewed by Cursor Bugbot for commit 503337f. Bugbot is set up for automated code reviews on this repo. Configure here.