close
Skip to content

[REPL_ENV] update RLM env to support recursion - #437

Merged
burtenshaw merged 39 commits into
huggingface:mainfrom
kashif:repl_env_update
Mar 20, 2026
Merged

[REPL_ENV] update RLM env to support recursion#437
burtenshaw merged 39 commits into
huggingface:mainfrom
kashif:repl_env_update

Conversation

@kashif

@kashif kashif commented Mar 15, 2026

Copy link
Copy Markdown
Collaborator

Summary

Added support for recursion for RLMs

Type of Change

  • Bug fix
  • New feature
  • Breaking change
  • Documentation
  • New environment
  • Refactoring

Alignment Checklist

Before submitting, verify:

  • I have read .claude/docs/PRINCIPLES.md and this PR aligns with our principles
  • I have checked .claude/docs/INVARIANTS.md and no invariants are violated
  • I have run /pre-submit-pr (or bash .claude/hooks/lint.sh and tests) and addressed all issues

RFC Status

  • Not required (bug fix, docs, minor refactoring)
  • RFC exists: #___
  • RFC needed (will create before merge)

Test Plan

Claude Code Review

@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Meta Open Source bot. label Mar 15, 2026
@greptile-apps

greptile-apps Bot commented Mar 15, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

This PR refactors the repl_env environment to support recursive child RLM runs by introducing three new modules: recursive_backends.py (shared-limit tree backend), recursive_controller.py (server-side composition), and runner.py (local orchestration loop). The unified REPLEnv is split into a pure remote EnvClient (client.py) and an explicit in-process helper (LocalREPLEnv in local.py), following the OpenEnv client/server separation principle.

Key issues found:

  • _apply_rubric method is undefined on REPLEnvironment but called at three sites in step() (lines 439, 452, 492). Every env.step() / env.execute() call — including through LocalRLMRunner and LocalREPLEnv — will raise AttributeError at runtime. The test suite is gated behind pytest.importorskip("smolagents"), so this is silently skipped in CI. (Flagged in prior thread — fix is to add def _apply_rubric(self, action, obs): return self.rubric(action, obs).)
  • REPLObservation in models.py does not declare done, reward, or metadata fields, yet these are used as constructor kwargs and dynamic assignments throughout repl_environment.py, and accessed in local.py._wrap_observation. Whether this works depends entirely on the opaque base Observation class — these fields should be declared explicitly.
  • Greedy regex in _default_answer (runner.py:247) uses FINAL\((.*)\) while _extract_final_answer correctly uses the lazy FINAL\((.*?)\). When a model response contains two FINAL() patterns, the greedy version spans between them, extracting the wrong content.
  • Architectural alignment: REPL_RLM_MAX_DEPTH defaults to 2 in app.py, enabling recursive child spawning on every server instance by default. Recursive behaviour allocates threads and creates LocalRLMRunner sub-processes during inference — consider defaulting to 1 and requiring explicit opt-in, or documenting the resource implications clearly.

What is fixed in this PR vs. prior threads:

  • The tree-global _children_spawned counter race condition is correctly addressed by moving the counter and lock into the shared BackendLimits dataclass.
  • The depth_changed false-positive in LocalREPLEnv is fixed by removing the pre-assignments from local.py.reset().
  • The per-child timeout is now cooperative (checked between runner iterations) rather than relying on ThreadPoolExecutor.result(timeout=...), eliminating the thread-blocking and concurrent.futures.TimeoutError class mismatch.

Confidence Score: 1/5

  • Not safe to merge — _apply_rubric is undefined and will crash every step() call at runtime.
  • The undefined _apply_rubric method means the entire environment is broken at runtime for any episode that calls step(). The test suite is silently skipped in CI (gated on smolagents), masking the failure. Additionally, REPLObservation does not explicitly declare done/reward/metadata fields, introducing fragile reliance on the base class. Until _apply_rubric is defined and the model fields are made explicit, the core execution path is non-functional.
  • envs/repl_env/server/repl_environment.py (missing _apply_rubric method), envs/repl_env/models.py (missing done/reward/metadata on REPLObservation), envs/repl_env/runner.py (greedy regex in _default_answer)

Important Files Changed

Filename Overview
envs/repl_env/recursive_backends.py New file implementing DirectLMBackend and LocalChildRLMBackend. The shared BackendLimits dataclass correctly uses a threading.Lock to guard the tree-global _children_spawned counter, resolving the earlier per-node race conditions. Cooperative timeout (per_child_timeout_s) is now propagated to child.run() rather than relying on ThreadPoolExecutor.result(timeout=...), which is a correct and documented trade-off.
envs/repl_env/runner.py New LocalRLMRunner orchestration loop. Two issues: (1) _default_answer uses a greedy regex FINAL\((.*)\) that can over-capture when a model response contains multiple FINAL() patterns — should use lazy (.*?) as _extract_final_answer does. (2) _default_backend_factory builds BackendLimits from self.max_depth instead of kwargs["max_depth"], creating an asymmetric API surface for custom backend factories.
envs/repl_env/local.py New LocalREPLEnv in-process helper. Cleanly wraps REPLEnvironment and no longer pre-assigns rlm_max_depth before calling reset() (fixing the false depth_changed detection from prior threads). _wrap_observation accesses obs.reward and obs.done which must be defined on the base Observation class — not visible in models.py.
envs/repl_env/server/repl_environment.py Critical: self._apply_rubric(action, obs) is called at lines 439, 452, and 492 but _apply_rubric is never defined on REPLEnvironment — every step() call will raise AttributeError. Additionally, the depth_changed rebuild path reuses _runtime_controller_chat_fn without re-validating credentials, and silent failure on missing huggingface_hub leaves LLM functions unconfigured with no user-facing warning (both noted in prior threads).
envs/repl_env/models.py REPLObservation does not declare done, reward, or metadata fields. These are used throughout repl_environment.py (constructor kwargs and dynamic attribute assignment) and accessed in local.py._wrap_observation. If the base Observation class does not provide them, construction and access will raise Pydantic ValidationError or AttributeError.
envs/repl_env/rubrics.py New rubric implementations (ExactMatchRubric, FuzzyMatchRubric, CustomMetricRubric, CodeExecutionRubric, REPLRubric). All correctly follow the Rubric interface with forward() and reset(). Well-structured and separates outcome vs. process rewards per RFC 004 design.
envs/repl_env/server/app.py Clean server factory. Correctly wires env vars to REPLEnvironment constructor. REPL_RLM_MAX_DEPTH defaults to 2, enabling recursive child runners by default on the server side — this is consistent with the RLM use case but means every server instance creates LocalChildRLMBackend unless explicitly disabled via env var.

Sequence Diagram

sequenceDiagram
    participant User
    participant LocalRLMRunner
    participant LocalREPLEnv
    participant REPLEnvironment
    participant Backend as LocalChildRLMBackend
    participant ChildRunner as LocalRLMRunner (child)

    User->>LocalRLMRunner: run(context, task_prompt)
    LocalRLMRunner->>Backend: backend_factory(llm_chat_fn, ...)
    Backend-->>LocalRLMRunner: backend instance (shared BackendLimits)

    LocalRLMRunner->>LocalREPLEnv: __init__(llm_query_fn, subcall_fn, ...)
    LocalRLMRunner->>REPLEnvironment: reset(context, task_prompt, max_iterations)
    REPLEnvironment-->>LocalRLMRunner: REPLObservation (done=False)

    loop up to max_iterations
        LocalRLMRunner->>LocalRLMRunner: _chat(messages, model)
        LocalRLMRunner->>LocalREPLEnv: execute(code)
        LocalREPLEnv->>REPLEnvironment: step(REPLAction)

        alt code calls rlm_query(prompt)
            REPLEnvironment->>Backend: recursive_query(prompt)
            Backend->>Backend: check _children_lock / _children_spawned
            Backend->>ChildRunner: runner_factory(..., depth=N+1)
            ChildRunner->>ChildRunner: run(prompt, prompt, timeout_s=per_child_timeout_s)
            ChildRunner-->>Backend: RLMRunResult.final_answer
            Backend-->>REPLEnvironment: child answer string
        end

        REPLEnvironment-->>LocalREPLEnv: REPLObservation
        LocalREPLEnv-->>LocalRLMRunner: StepResult

        alt done == True
            LocalRLMRunner-->>User: RLMRunResult
        end
    end

    LocalRLMRunner->>LocalRLMRunner: _default_answer(messages) [fallback]
    LocalRLMRunner-->>User: RLMRunResult
Loading

Comments Outside Diff (1)

  1. envs/repl_env/models.py, line 65-86 (link)

    P2 REPLObservation has no reward, done, or metadata fields — LocalREPLEnv._wrap_observation will raise AttributeError on every reset call

    local.py._wrap_observation unconditionally accesses obs.reward and obs.done:

    return StepResult(
        observation=obs,
        reward=obs.reward,   # ← not defined on REPLObservation
        done=obs.done,       # ← not defined on REPLObservation
    )

    REPLEnvironment.reset() also constructs REPLObservation(done=False, metadata={…}) — fields not present in the Pydantic model definition. With Pydantic v2, passing unknown keyword arguments to a model constructor raises ValidationError by default unless the base Observation class defines them or sets model_config = ConfigDict(extra="allow").

    Even if the base class provides done and metadata, reward is never set on the reset-path observation (only step() calls _apply_rubric), so accessing obs.reward in _wrap_observation after a reset() would return None only if Observation defines it as Optional[float] = None — which is not visible in this model file.

    These fields (done, reward, metadata) should be explicitly declared on REPLObservation with sensible defaults to avoid relying on opaque base-class definitions:

    done: bool = Field(default=False, description="Whether the episode is done")
    reward: Optional[float] = Field(default=None, description="Reward for this step")
    metadata: Dict[str, Any] = Field(default_factory=dict, description="Step metadata")
    Prompt To Fix With AI
    This is a comment left during a code review.
    Path: envs/repl_env/models.py
    Line: 65-86
    
    Comment:
    **`REPLObservation` has no `reward`, `done`, or `metadata` fields — `LocalREPLEnv._wrap_observation` will raise `AttributeError` on every reset call**
    
    `local.py._wrap_observation` unconditionally accesses `obs.reward` and `obs.done`:
    
    ```python
    return StepResult(
        observation=obs,
        reward=obs.reward,   # ← not defined on REPLObservation
        done=obs.done,       # ← not defined on REPLObservation
    )
    ```
    
    `REPLEnvironment.reset()` also constructs `REPLObservation(done=False, metadata={…})` — fields not present in the Pydantic model definition. With Pydantic v2, passing unknown keyword arguments to a model constructor raises `ValidationError` by default unless the base `Observation` class defines them or sets `model_config = ConfigDict(extra="allow")`.
    
    Even if the base class provides `done` and `metadata`, `reward` is never set on the reset-path observation (only `step()` calls `_apply_rubric`), so accessing `obs.reward` in `_wrap_observation` after a `reset()` would return `None` only if `Observation` defines it as `Optional[float] = None` — which is not visible in this model file.
    
    These fields (`done`, `reward`, `metadata`) should be explicitly declared on `REPLObservation` with sensible defaults to avoid relying on opaque base-class definitions:
    
    ```python
    done: bool = Field(default=False, description="Whether the episode is done")
    reward: Optional[float] = Field(default=None, description="Reward for this step")
    metadata: Dict[str, Any] = Field(default_factory=dict, description="Step metadata")
    ```
    
    How can I resolve this? If you propose a fix, please make it concise.
Prompt To Fix All With AI
This is a comment left during a code review.
Path: envs/repl_env/runner.py
Line: 247-248

Comment:
**Greedy regex in `_default_answer` will over-capture when response contains multiple `FINAL(...)` calls**

`_extract_final_answer` (in `repl_environment.py`) uses the lazy quantifier `FINAL\((.*?)\)` to avoid over-capturing. `_default_answer` uses the greedy `FINAL\((.*)\)` with `re.DOTALL`. When a model returns:

```
I think the answer is FINAL(42). Here's why: …
```

this works fine. But when it returns two patterns:

```
My first thought was FINAL(wrong) but actually FINAL(correct)
```

the greedy `.+` spans from the first `FINAL(` all the way to the last `)`, returning `wrong) but actually FINAL(correct` as the captured group — completely wrong.

Use the same lazy pattern as `_extract_final_answer`:

```suggestion
            match = re.search(r"FINAL\((.*?)\)", response, re.DOTALL)
```

How can I resolve this? If you propose a fix, please make it concise.

---

This is a comment left during a code review.
Path: envs/repl_env/models.py
Line: 65-86

Comment:
**`REPLObservation` has no `reward`, `done`, or `metadata` fields — `LocalREPLEnv._wrap_observation` will raise `AttributeError` on every reset call**

`local.py._wrap_observation` unconditionally accesses `obs.reward` and `obs.done`:

```python
return StepResult(
    observation=obs,
    reward=obs.reward,   # ← not defined on REPLObservation
    done=obs.done,       # ← not defined on REPLObservation
)
```

`REPLEnvironment.reset()` also constructs `REPLObservation(done=False, metadata={…})` — fields not present in the Pydantic model definition. With Pydantic v2, passing unknown keyword arguments to a model constructor raises `ValidationError` by default unless the base `Observation` class defines them or sets `model_config = ConfigDict(extra="allow")`.

Even if the base class provides `done` and `metadata`, `reward` is never set on the reset-path observation (only `step()` calls `_apply_rubric`), so accessing `obs.reward` in `_wrap_observation` after a `reset()` would return `None` only if `Observation` defines it as `Optional[float] = None` — which is not visible in this model file.

These fields (`done`, `reward`, `metadata`) should be explicitly declared on `REPLObservation` with sensible defaults to avoid relying on opaque base-class definitions:

```python
done: bool = Field(default=False, description="Whether the episode is done")
reward: Optional[float] = Field(default=None, description="Reward for this step")
metadata: Dict[str, Any] = Field(default_factory=dict, description="Step metadata")
```

How can I resolve this? If you propose a fix, please make it concise.

---

This is a comment left during a code review.
Path: envs/repl_env/runner.py
Line: 86-107

Comment:
**`_default_backend_factory` ignores its `kwargs["max_depth"]` in favour of `self.max_depth`**

The factory receives `max_depth` via `kwargs` (forwarded by `run()` at line 121), but builds the `BackendLimits` from `self.max_depth` instead:

```python
limits = BackendLimits(
    max_depth=self.max_depth,   # ← uses attribute, not kwargs["max_depth"]
    ...
)
```

Today this is harmless because `run()` always passes `max_depth=self.max_depth`, so both values are identical. However, the asymmetry creates a misleading API: a custom `backend_factory` modelled on this default would capture `self.max_depth` from the wrong scope (the runner's attribute) rather than the forwarded kwarg. Consider using `kwargs["max_depth"]` for consistency:

```python
limits = BackendLimits(
    max_depth=kwargs["max_depth"],
    ...
)
```

How can I resolve this? If you propose a fix, please make it concise.

Last reviewed commit: "Merge branch 'main' ..."

Comment thread envs/repl_env/client.py Outdated
Comment thread examples/repl_oolong_simple.py Outdated
@kashif kashif changed the title [REPL_ENV] update RLM env [REPL_ENV] update RLM env to support recursion Mar 16, 2026
@burtenshaw

Copy link
Copy Markdown
Collaborator

@greptile-apps

Comment thread envs/repl_env/recursive_backends.py Outdated
Comment thread envs/repl_env/server/repl_environment.py
Comment thread envs/repl_env/recursive_controller.py Outdated
Comment thread envs/repl_env/recursive_backends.py Outdated
Comment thread envs/repl_env/recursive_backends.py Outdated
on_subcall_complete=self.on_subcall_complete,
)
if self.limits.per_child_timeout_s is None:
result = child.run(prompt, prompt, model=model)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same string used as both context and task_prompt

child.run(prompt, prompt, model=model) passes the user's prompt as both the context variable and the task_prompt. This means when a child RLM runner is spawned, it loads the question text itself as the context object inside the REPL. While functional (the child can still answer by reasoning on it), it is likely to confuse models: the system prompt tells them to "check the content of the context variable" which contains their task, not a separate document.

Consider whether a child should receive an empty context ("") and the prompt as only the task, or whether this is intentional for self-referential tasks:

Suggested change
result = child.run(prompt, prompt, model=model)
result = child.run("", prompt, model=model)
Prompt To Fix With AI
This is a comment left during a code review.
Path: envs/repl_env/recursive_backends.py
Line: 179

Comment:
**Same string used as both `context` and `task_prompt`**

`child.run(prompt, prompt, model=model)` passes the user's `prompt` as both the `context` variable and the `task_prompt`. This means when a child RLM runner is spawned, it loads the question text itself as the `context` object inside the REPL. While functional (the child can still answer by reasoning on it), it is likely to confuse models: the system prompt tells them to "check the content of the `context` variable" which contains their task, not a separate document.

Consider whether a child should receive an empty context (`""`) and the prompt as only the task, or whether this is intentional for self-referential tasks:

```suggestion
            result = child.run("", prompt, model=model)
```

How can I resolve this? If you propose a fix, please make it concise.

Comment thread envs/repl_env/recursive_controller.py
@kashif

kashif commented Mar 16, 2026

Copy link
Copy Markdown
Collaborator Author

@greptile-apps

Comment thread envs/repl_env/recursive_backends.py Outdated
Comment thread envs/repl_env/runner.py
kashif and others added 2 commits March 16, 2026 13:47
Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
@kashif

kashif commented Mar 16, 2026

Copy link
Copy Markdown
Collaborator Author

@greptile-apps

Comment thread envs/repl_env/recursive_backends.py Outdated
Comment thread envs/repl_env/recursive_backends.py Outdated
Comment thread envs/repl_env/runner.py
kashif and others added 2 commits March 16, 2026 14:37
Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
@kashif

kashif commented Mar 16, 2026

Copy link
Copy Markdown
Collaborator Author

@greptile-apps

Comment on lines +239 to +248
self.rlm_max_depth = runtime_rlm_max_depth
self.rlm_max_iterations = runtime_rlm_max_iterations

# Create LLM functions if not already provided at init.
# Match example/client behavior by allowing anonymous routed inference too.
if not self.llm_query_fn:
effective_token = hf_token or os.environ.get("HF_TOKEN")
if effective_token:
self._create_llm_functions(effective_token, llm_model)
effective_token = (
hf_token if hf_token is not None else os.environ.get("HF_TOKEN")
)
self._create_llm_functions(effective_token, llm_model)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Recursive controller not rebuilt when rlm_max_depth changes on re-reset

self.rlm_max_depth is updated on lines 239–240 when a caller passes new rlm_max_depth/rlm_max_iterations kwargs. However, the guard on line 244 — if not self.llm_query_fn — is already True after the first successful reset (because self.llm_query_fn was set by _create_llm_functions). This means any updated depth/iteration values are written to the instance attributes but the underlying RecursiveController (and its LocalChildRLMBackend or DirectLMBackend) is never re-created to reflect them.

Concrete failure scenario:

  1. Client resets with no kwargs → controller is built with rlm_max_depth=1 (default), rlm_query_fn=None.
  2. Client resets again with rlm_max_depth=2self.rlm_max_depth becomes 2, but the old depth-1 controller is still active, so rlm_query remains absent from the REPL namespace.

The fix is to also trigger _create_llm_functions when the runtime depth or iteration config changes:

depth_changed = (
    runtime_rlm_max_depth != self.rlm_max_depth
    or runtime_rlm_max_iterations != self.rlm_max_iterations
)
if not self.llm_query_fn or depth_changed:
    ...
    self._create_llm_functions(effective_token, llm_model)

(Store the depth/iterations before updating self so the comparison is meaningful.)

Prompt To Fix With AI
This is a comment left during a code review.
Path: envs/repl_env/server/repl_environment.py
Line: 239-248

Comment:
**Recursive controller not rebuilt when `rlm_max_depth` changes on re-reset**

`self.rlm_max_depth` is updated on lines 239–240 when a caller passes new `rlm_max_depth`/`rlm_max_iterations` kwargs. However, the guard on line 244 — `if not self.llm_query_fn` — is already `True` after the *first* successful reset (because `self.llm_query_fn` was set by `_create_llm_functions`). This means any updated depth/iteration values are written to the instance attributes but the underlying `RecursiveController` (and its `LocalChildRLMBackend` or `DirectLMBackend`) is **never re-created** to reflect them.

Concrete failure scenario:
1. Client resets with no kwargs → controller is built with `rlm_max_depth=1` (default), `rlm_query_fn=None`.
2. Client resets again with `rlm_max_depth=2``self.rlm_max_depth` becomes `2`, but the old depth-1 controller is still active, so `rlm_query` remains absent from the REPL namespace.

The fix is to also trigger `_create_llm_functions` when the runtime depth or iteration config changes:

```python
depth_changed = (
    runtime_rlm_max_depth != self.rlm_max_depth
    or runtime_rlm_max_iterations != self.rlm_max_iterations
)
if not self.llm_query_fn or depth_changed:
    ...
    self._create_llm_functions(effective_token, llm_model)
```

(Store the depth/iterations *before* updating `self` so the comparison is meaningful.)

How can I resolve this? If you propose a fix, please make it concise.

Comment on lines +176 to +180
"""Create LLM/subcall functions dynamically using client-provided token."""
try:
chat_fn = self._build_hf_chat_fn(hf_token, llm_model)
except RuntimeError:
return

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Silent failure when huggingface_hub is not installed

_create_llm_functions catches RuntimeError from _build_hf_chat_fn (raised when huggingface_hub is missing) and returns silently. A caller that passed hf_token expecting LLM support gets none, with no log line or warning. The LLM functions won't be injected into the REPL namespace, and the model will see llm_query/rlm_query as undefined — an inscrutable failure.

At minimum, emit a warning:

    except RuntimeError as exc:
        import warnings
        warnings.warn(
            f"[repl_env] Could not initialize LLM functions: {exc}. "
            "llm_query/rlm_query will not be available in the REPL.",
            RuntimeWarning,
            stacklevel=2,
        )
        return
Prompt To Fix With AI
This is a comment left during a code review.
Path: envs/repl_env/server/repl_environment.py
Line: 176-180

Comment:
**Silent failure when `huggingface_hub` is not installed**

`_create_llm_functions` catches `RuntimeError` from `_build_hf_chat_fn` (raised when `huggingface_hub` is missing) and returns silently. A caller that passed `hf_token` expecting LLM support gets none, with no log line or warning. The LLM functions won't be injected into the REPL namespace, and the model will see `llm_query`/`rlm_query` as undefined — an inscrutable failure.

At minimum, emit a warning:

```python
    except RuntimeError as exc:
        import warnings
        warnings.warn(
            f"[repl_env] Could not initialize LLM functions: {exc}. "
            "llm_query/rlm_query will not be available in the REPL.",
            RuntimeWarning,
            stacklevel=2,
        )
        return
```

How can I resolve this? If you propose a fix, please make it concise.

Comment on lines 244 to +248
if not self.llm_query_fn:
effective_token = hf_token or os.environ.get("HF_TOKEN")
if effective_token:
self._create_llm_functions(effective_token, llm_model)
effective_token = (
hf_token if hf_token is not None else os.environ.get("HF_TOKEN")
)
self._create_llm_functions(effective_token, llm_model)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LLM client created even when no auth token is available

_create_llm_functions is called whenever self.llm_query_fn is unset, including when effective_token is None (no token was supplied by the client or server environment). _build_hf_chat_fn accepts None for hf_token despite the str annotation, passes it to InferenceClient, and the resulting client object silently accepts all construction calls.

For gated models (including the server default), every inference call then fails at runtime with an authentication error. No diagnostic appears at reset() time — the failure surfaces mid-episode as an opaque error string returned from the injected REPL function.

The guard on this line should also require that effective_token is non-None before constructing the client, so that the fast-fail happens at initialization rather than during inference. Additionally, the hf_token parameter type annotation in _build_hf_chat_fn should be Optional[str].

Prompt To Fix With AI
This is a comment left during a code review.
Path: envs/repl_env/server/repl_environment.py
Line: 244-248

Comment:
**LLM client created even when no auth token is available**

`_create_llm_functions` is called whenever `self.llm_query_fn` is unset, including when `effective_token` is `None` (no token was supplied by the client or server environment). `_build_hf_chat_fn` accepts `None` for `hf_token` despite the `str` annotation, passes it to `InferenceClient`, and the resulting client object silently accepts all construction calls.

For gated models (including the server default), every inference call then fails at runtime with an authentication error. No diagnostic appears at `reset()` time — the failure surfaces mid-episode as an opaque error string returned from the injected REPL function.

The guard on this line should also require that `effective_token` is non-`None` before constructing the client, so that the fast-fail happens at initialization rather than during inference. Additionally, the `hf_token` parameter type annotation in `_build_hf_chat_fn` should be `Optional[str]`.

How can I resolve this? If you propose a fix, please make it concise.

@kashif

kashif commented Mar 17, 2026

Copy link
Copy Markdown
Collaborator Author

@greptile-apps

Comment thread envs/repl_env/local.py
Comment on lines +76 to +90
self._env.max_iterations = max_iterations
if rlm_max_depth is not None:
self._env.rlm_max_depth = rlm_max_depth
if rlm_max_iterations is not None:
self._env.rlm_max_iterations = rlm_max_iterations
reset_kwargs = {}
if rlm_max_depth is not None:
reset_kwargs["rlm_max_depth"] = rlm_max_depth
if rlm_max_iterations is not None:
reset_kwargs["rlm_max_iterations"] = rlm_max_iterations
if expected_answer is not None:
reset_kwargs["expected_answer"] = expected_answer
obs = self._env.reset(
seed=seed,
episode_id=episode_id,

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 depth_changed detection always evaluates false via LocalREPLEnv

LocalREPLEnv.reset() writes self._env.rlm_max_depth = rlm_max_depth on lines 77–78 before calling self._env.reset(...). Inside REPLEnvironment.reset(), the depth-change guard compares runtime_rlm_max_depth (taken from the kwarg) against self.rlm_max_depth (already updated by the pre-assignment). The comparison is therefore new_depth != new_depth, which is always False.

Concrete failure scenario:

  1. LocalREPLEnv() is constructed without llm_query_fn. The first reset(rlm_max_depth=1) builds a depth-1 controller.
  2. reset(rlm_max_depth=3) is called next — LocalREPLEnv pre-sets self._env.rlm_max_depth = 3, then calls self._env.reset(…, rlm_max_depth=3). Inside REPLEnvironment.reset(), depth_changed evaluates to 3 != 3 → False, so the old depth-1 controller is silently reused.

Fix: remove the redundant pre-assignments (lines 77–80) from LocalREPLEnv.reset() and rely solely on the kwarg-based path in REPLEnvironment.reset(), which was specifically designed to detect and handle the change. The values will still be updated correctly via the depth_changed path in REPLEnvironment.reset().

Prompt To Fix With AI
This is a comment left during a code review.
Path: envs/repl_env/local.py
Line: 76-90

Comment:
**`depth_changed` detection always evaluates false via `LocalREPLEnv`**

`LocalREPLEnv.reset()` writes `self._env.rlm_max_depth = rlm_max_depth` on lines 77–78 **before** calling `self._env.reset(...)`. Inside `REPLEnvironment.reset()`, the depth-change guard compares `runtime_rlm_max_depth` (taken from the kwarg) against `self.rlm_max_depth` (already updated by the pre-assignment). The comparison is therefore `new_depth != new_depth`, which is always `False`.

Concrete failure scenario:
1. `LocalREPLEnv()` is constructed without `llm_query_fn`. The first `reset(rlm_max_depth=1)` builds a depth-1 controller.
2. `reset(rlm_max_depth=3)` is called next — `LocalREPLEnv` pre-sets `self._env.rlm_max_depth = 3`, then calls `self._env.reset(…, rlm_max_depth=3)`. Inside `REPLEnvironment.reset()`, `depth_changed` evaluates to `3 != 3 → False`, so the old depth-1 controller is silently reused.

**Fix**: remove the redundant pre-assignments (lines 77–80) from `LocalREPLEnv.reset()` and rely solely on the kwarg-based path in `REPLEnvironment.reset()`, which was specifically designed to detect and handle the change. The values will still be updated correctly via the `depth_changed` path in `REPLEnvironment.reset()`.

How can I resolve this? If you propose a fix, please make it concise.

Comment on lines +249 to 256
# Token resolution: explicit hf_token > HF_TOKEN env var > cached HF login.
if not self.llm_query_fn or depth_changed:
effective_token = (
hf_token if hf_token is not None else os.environ.get("HF_TOKEN")
)
self._create_llm_functions(effective_token, llm_model)

# Initialize state

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Depth-change rebuild can clobber a valid LLM setup with a no-credential client

The original code only called _create_llm_functions when a valid API credential was available. This PR removed that guard when rebuilding due to depth_changed.

The consequence: if a caller provides credentials on the first reset() (valid controller is built), then calls reset(rlm_max_depth=3) on the next episode without re-providing credentials — depth_changed=True and no credential is resolved, so _create_llm_functions(None, …) is invoked. The old valid controller is replaced by one backed by an InferenceClient constructed without auth. All subsequent llm_query / rlm_query calls then fail at inference time with auth errors, with no diagnostic at reset time.

The fix is to restore the credential guard so _create_llm_functions is only called when a valid credential is present, regardless of whether the trigger is a missing llm_query_fn or the depth_changed flag.

Prompt To Fix With AI
This is a comment left during a code review.
Path: envs/repl_env/server/repl_environment.py
Line: 249-256

Comment:
**Depth-change rebuild can clobber a valid LLM setup with a no-credential client**

The original code only called `_create_llm_functions` when a valid API credential was available. This PR removed that guard when rebuilding due to `depth_changed`.

The consequence: if a caller provides credentials on the first `reset()` (valid controller is built), then calls `reset(rlm_max_depth=3)` on the next episode without re-providing credentials — `depth_changed=True` and no credential is resolved, so `_create_llm_functions(None, …)` is invoked. The old valid controller is **replaced** by one backed by an `InferenceClient` constructed without auth. All subsequent `llm_query` / `rlm_query` calls then fail at inference time with auth errors, with no diagnostic at reset time.

The fix is to restore the credential guard so `_create_llm_functions` is only called when a valid credential is present, regardless of whether the trigger is a missing `llm_query_fn` or the `depth_changed` flag.

How can I resolve this? If you propose a fix, please make it concise.

@kashif

kashif commented Mar 17, 2026

Copy link
Copy Markdown
Collaborator Author

@greptile-apps

message="Final answer submitted.",
reward=self.reward_on_success,
)
obs.reward = self._apply_rubric(action, obs)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P0 _apply_rubric method undefined — every step() call will raise AttributeError

self._apply_rubric(action, obs) is called on lines 439, 452, and 492, but the method is never defined anywhere in REPLEnvironment. At runtime, every call to env.step() — and therefore every call to env.execute() via LocalREPLEnv or LocalRLMRunner — will raise:

AttributeError: 'REPLEnvironment' object has no attribute '_apply_rubric'

The rubric is stored at self.rubric (a REPLRubric instance). The method body should be:

def _apply_rubric(self, action: REPLAction, obs: REPLObservation) -> float:
    return self.rubric(action, obs)

Why this wasn't caught: the entire test file starts with pytest.importorskip("smolagents"). If smolagents is absent in the CI environment, every test — including test_rubric_reward_on_success, test_rubric_reward_on_error, and the full TestLocalRLMRunner suite — is silently skipped. The pre-submit checks would show "0 failures" while the environment is completely broken for any caller that actually uses step().

Prompt To Fix With AI
This is a comment left during a code review.
Path: envs/repl_env/server/repl_environment.py
Line: 439

Comment:
**`_apply_rubric` method undefined — every `step()` call will raise `AttributeError`**

`self._apply_rubric(action, obs)` is called on lines 439, 452, and 492, but the method is never defined anywhere in `REPLEnvironment`. At runtime, every call to `env.step()` — and therefore every call to `env.execute()` via `LocalREPLEnv` or `LocalRLMRunner` — will raise:

```
AttributeError: 'REPLEnvironment' object has no attribute '_apply_rubric'
```

The rubric is stored at `self.rubric` (a `REPLRubric` instance). The method body should be:

```python
def _apply_rubric(self, action: REPLAction, obs: REPLObservation) -> float:
    return self.rubric(action, obs)
```

**Why this wasn't caught:** the entire test file starts with `pytest.importorskip("smolagents")`. If `smolagents` is absent in the CI environment, every test — including `test_rubric_reward_on_success`, `test_rubric_reward_on_error`, and the full `TestLocalRLMRunner` suite — is silently skipped. The pre-submit checks would show "0 failures" while the environment is completely broken for any caller that actually uses `step()`.

How can I resolve this? If you propose a fix, please make it concise.

@kashif

kashif commented Mar 18, 2026

Copy link
Copy Markdown
Collaborator Author

@greptile-apps

Comment thread envs/repl_env/runner.py Outdated
@burtenshaw
burtenshaw merged commit 7f86a13 into huggingface:main Mar 20, 2026
8 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Meta Open Source bot.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants