Merge pull request #165 from alexzhang13/training_harness

Add prime-rl-based RLM training harness and update the local REPL implementation.

Author: alexzhang13

README.md

@@ -32,6 +32,8 @@ Recursive Language Models (RLMs) are a task-agnostic inference paradigm for lang
 
 This repository provides an extensible inference engine for using RLMs around standard API-based and local LLMs. The initial experiments and idea were proposed in a [blogpost](https://alexzhang13.github.io/blog/2025/rlm/) in 2025, with expanded results in an [arXiv preprint](https://arxiv.org/abs/2512.24601).
 
+We now also include a [verifiers](https://github.com/PrimeIntellect-ai/verifiers) training environment based on Prime Intellect's [prime-rl](https://github.com/PrimeIntellect-ai/prime-rl) in the `training/` folder. Train your own RLMs, which directly can be plugged into our inference engine!
+
 > [!NOTE]
 > This repository contains inference code for RLMs with support for various sandbox environments. Open-source contributions are welcome. This repository is maintained by the authors of the paper from the MIT OASYS lab.
 
@@ -123,6 +125,11 @@ export PRIME_API_KEY=...
 ### Model Providers
 We currently support most major clients (OpenAI, Anthropic), as well as the router platforms (OpenRouter, Portkey). For local models, we recommend using vLLM (which interfaces with the [OpenAI client](https://github.com/alexzhang13/rlm/blob/main/rlm/clients/openai.py)). To view or add support for more clients, start by looking at [`rlm/clients/`](https://github.com/alexzhang13/rlm/tree/main/rlm/clients).
 
+## Training
+We provide a simple RL training harness for training RLMs used in this repo (specifically the `local` REPL). The implementation uses no sandboxes for simplicity and slots easily your use case, but an ideal setup would use sandboxes for safety. Training logic is isolated to the [`training/`](https://github.com/alexzhang13/rlm/tree/main/training) folder, which exposes `rlm.RLM` as a [`verifiers`](https://github.com/willccbb/verifiers) `Environment` and plugs straight into [`prime-rl`](https://github.com/PrimeIntellect-ai/prime-rl). See the [training README](https://github.com/alexzhang13/rlm/tree/main/training#readme) for the launch command. The harness uses subprocess-isolated local REPL execution (no cloud sandboxes), matching the `local` environment above.
+
+A worked example with an example `.toml` lives in [`training/environments/oolong/`](https://github.com/alexzhang13/rlm/tree/main/training/environments/oolong) (OOLONG long-context QA). New training environments can be added the same way — author a `verifiers` env that wraps your task (see the [verifiers docs](https://verifiers.readthedocs.io/)), then reference it from a config.
+
 ## Relevant Reading
 * **[Dec '25]** [Recursive Language Models arXiv](https://arxiv.org/abs/2512.24601)
 * **[Oct '25]** [Recursive Language Models Blogpost](https://alexzhang13.github.io/blog/2025/rlm/)

rlm/clients/base_lm.py

@@ -13,9 +13,19 @@ class BaseLM(ABC):
     does so in a model-agnostic way, so this class provides a base interface for all language models.
     """
 
-    def __init__(self, model_name: str, timeout: float = DEFAULT_TIMEOUT, **kwargs):
+    def __init__(
+        self,
+        model_name: str,
+        timeout: float = DEFAULT_TIMEOUT,
+        sampling_args: dict[str, Any] | None = None,
+        **kwargs,
+    ):
         self.model_name = model_name
         self.timeout = timeout
+        # Sampling args forwarded to the underlying completion API
+        # (e.g. temperature, top_p, max_tokens, seed). Forwarded by
+        # subclasses as **self.sampling_args.
+        self.sampling_args: dict[str, Any] = dict(sampling_args or {})
         self.kwargs = kwargs
 
     @abstractmethod

rlm/clients/openai.py

@@ -18,6 +18,30 @@
 DEFAULT_PRIME_INTELLECT_BASE_URL = "https://api.pinference.ai/api/v1/"
 
 
+def _normalize_sampling_args(sampling_args: dict[str, Any]) -> dict[str, Any]:
+    """Match the rename done by verifiers' OpenAIChatCompletionsClient so the
+    same sampling_args dict produces byte-equivalent chat.completions.create
+    calls in both harnesses. Pops ``extra_body`` so the caller can merge it
+    with its own ``extra_body`` rather than passing it twice (TypeError).
+    """
+    args = dict(sampling_args or {})
+    if "max_tokens" in args:
+        args["max_completion_tokens"] = args.pop("max_tokens")
+    args.pop("extra_body", None)
+    return {k: v for k, v in args.items() if v is not None}
+
+
+def _merge_extra_body(
+    hardcoded: dict[str, Any], sampling_args: dict[str, Any] | None
+) -> dict[str, Any]:
+    """Merge an ``extra_body`` from sampling_args into the hardcoded extra_body."""
+    merged = dict(hardcoded or {})
+    user = (sampling_args or {}).get("extra_body")
+    if user:
+        merged.update(user)
+    return merged
+
+
 class OpenAIClient(BaseLM):
     """
     LM Client for running models with the OpenAI API. Works with vLLM as well.
@@ -32,9 +56,10 @@ def __init__(
         api_key: str | None = None,
         model_name: str | None = None,
         base_url: str | None = None,
+        sampling_args: dict[str, Any] | None = None,
         **kwargs,
     ):
-        super().__init__(model_name=model_name, **kwargs)
+        super().__init__(model_name=model_name, sampling_args=sampling_args, **kwargs)
 
         if api_key is None:
             if base_url == "https://api.openai.com/v1" or base_url is None:
@@ -78,12 +103,16 @@ def completion(self, prompt: str | list[dict[str, Any]], model: str | None = Non
         if not model:
             raise ValueError("Model name is required for OpenAI client.")
 
-        extra_body = {}
+        extra_body: dict[str, Any] = {}
         if self.client.base_url == DEFAULT_PRIME_INTELLECT_BASE_URL:
             extra_body["usage"] = {"include": True}
+        extra_body = _merge_extra_body(extra_body, self.sampling_args)
 
         response = self.client.chat.completions.create(
-            model=model, messages=messages, extra_body=extra_body
+            model=model,
+            messages=messages,
+            extra_body=extra_body,
+            **_normalize_sampling_args(self.sampling_args),
         )
         self._track_cost(response, model)
         return response.choices[0].message.content
@@ -102,12 +131,16 @@ async def acompletion(
         if not model:
             raise ValueError("Model name is required for OpenAI client.")
 
-        extra_body = {}
+        extra_body: dict[str, Any] = {}
         if self.client.base_url == DEFAULT_PRIME_INTELLECT_BASE_URL:
             extra_body["usage"] = {"include": True}
+        extra_body = _merge_extra_body(extra_body, self.sampling_args)
 
         response = await self.async_client.chat.completions.create(
-            model=model, messages=messages, extra_body=extra_body
+            model=model,
+            messages=messages,
+            extra_body=extra_body,
+            **_normalize_sampling_args(self.sampling_args),
         )
         self._track_cost(response, model)
         return response.choices[0].message.content

rlm/core/rlm.py

@@ -74,6 +74,10 @@ def __init__(
         on_subcall_complete: Callable[[int, str, float, str | None], None] | None = None,
         on_iteration_start: Callable[[int, int], None] | None = None,
         on_iteration_complete: Callable[[int, int, float], None] | None = None,
+        sampling_args: dict[str, Any] | None = None,
+        sub_sampling_args: dict[str, Any] | None = None,
+        orchestrator: bool = True,
+        user_prologue: str | None = None,
     ):
         """
         Args:
@@ -109,6 +113,30 @@ def __init__(
             on_iteration_start: Callback fired when an iteration starts. Args: (depth, iteration_num).
             on_iteration_complete: Callback fired when an iteration completes. Args: (depth, iteration_num, duration).
         """
+        # Sampling args plumbed into backend_kwargs / other_backend_kwargs
+        # before the clients are constructed, so they reach the chat-completions
+        # call (e.g. temperature, top_p, max_tokens, seed). ``sampling_args``
+        # applies to the root model (depth=0); ``sub_sampling_args`` to
+        # depth=1 sub-LLM calls. If ``sub_sampling_args`` is set without an
+        # ``other_backends``, we mirror the root backend so depth=1 routes
+        # through a separate client with its own sampling args.
+        if sampling_args is not None:
+            backend_kwargs = dict(backend_kwargs or {})
+            existing = dict(backend_kwargs.get("sampling_args") or {})
+            existing.update(sampling_args)
+            backend_kwargs["sampling_args"] = existing
+        if sub_sampling_args is not None:
+            if other_backends is None:
+                other_backends = [backend]
+                other_backend_kwargs = [dict(backend_kwargs or {})]
+            else:
+                other_backend_kwargs = [dict(kw or {}) for kw in (other_backend_kwargs or [{}])]
+            first = dict(other_backend_kwargs[0])
+            existing = dict(first.get("sampling_args") or {})
+            existing.update(sub_sampling_args)
+            first["sampling_args"] = existing
+            other_backend_kwargs[0] = first
+
         # Store config for spawning per-completion
         self.backend = backend
         self.backend_kwargs = backend_kwargs
@@ -144,6 +172,12 @@ def __init__(
         self.max_tokens = max_tokens
         self.max_errors = max_errors
         self.system_prompt = custom_system_prompt if custom_system_prompt else RLM_SYSTEM_PROMPT
+        self.orchestrator = orchestrator
+        # Optional user-prologue message inserted between the metadata user
+        # message and the iter-0 turn prompt. Mirrors RLMTrainEnv's
+        # ``user_prologue`` so canonical inference can match envs that
+        # depend on a task-specific tips message (e.g. BC+).
+        self.user_prologue = user_prologue
         self.logger = logger
         self.verbose = VerbosePrinter(enabled=verbose)
 
@@ -261,7 +295,11 @@ def _spawn_completion_context(self, prompt: str | dict[str, Any]):
             if not self.persistent and hasattr(environment, "cleanup"):
                 environment.cleanup()
 
-    def _setup_prompt(self, prompt: str | dict[str, Any]) -> list[dict[str, Any]]:
+    def _setup_prompt(
+        self,
+        prompt: str | dict[str, Any],
+        root_prompt: str | None = None,
+    ) -> list[dict[str, Any]]:
         """
         Setup the system prompt for the RLM. Also include metadata about the prompt and build
         up the initial message history.
@@ -271,7 +309,11 @@ def _setup_prompt(self, prompt: str | dict[str, Any]) -> list[dict[str, Any]]:
             system_prompt=self.system_prompt,
             query_metadata=metadata,
             custom_tools=self.custom_tools,
+            root_prompt=root_prompt,
+            orchestrator=self.orchestrator,
         )
+        if self.user_prologue:
+            message_history.append({"role": "user", "content": self.user_prologue})
         if self.compaction:
             message_history[0]["content"] += (
                 "\n\nThe full conversation history (trajectory segments and any summaries) "
@@ -310,7 +352,7 @@ def completion(
             self.logger.clear_iterations()
 
         with self._spawn_completion_context(prompt) as (lm_handler, environment):
-            message_history = self._setup_prompt(prompt)
+            message_history = self._setup_prompt(prompt, root_prompt=root_prompt)
 
             compaction_count = 0
             try:
@@ -333,7 +375,6 @@ def completion(
                                 lm_handler, environment, message_history, compaction_count
                             )
 
-                    # Current prompt = message history + additional prompt suffix
                     context_count = (
                         environment.get_context_count()
                         if isinstance(environment, SupportsPersistence)
@@ -344,12 +385,22 @@ def completion(
                         if isinstance(environment, SupportsPersistence)
                         else 0
                     )
-                    current_prompt = message_history + [
-                        build_user_prompt(root_prompt, i, context_count, history_count)
-                    ]
+                    # Fully prefixed trajectory: persist the per-turn user prompt
+                    # into message_history so the model sees a single continuous
+                    # [system, metadata, user_0, assistant_0, repl_0, user_1, ...]
+                    # chain across turns.
+                    message_history.append(
+                        build_user_prompt(
+                            root_prompt,
+                            i,
+                            context_count,
+                            history_count,
+                            max_iterations=self.max_iterations,
+                        )
+                    )
 
                     iteration: RLMIteration = self._completion_turn(
-                        prompt=current_prompt,
+                        prompt=message_history,
                         lm_handler=lm_handler,
                         environment=environment,
                     )

rlm/utils/parsing.py

@@ -30,30 +30,42 @@ def format_iteration(
     the prompt of the LM in the next iteration. We also truncate code execution results
     that exceed the max_character_length.
 
+    Each iteration produces exactly two messages in history: one assistant
+    turn containing the model's response (with any ```repl``` blocks
+    embedded), followed by a single user message that concatenates the
+    outputs of all executed code blocks in that turn. This keeps the
+    per-turn shape assistant-then-user even when the model emits several
+    blocks in one response, and avoids redundantly echoing the code
+    (which is already in the assistant message) back in the user reply.
+    Each block's output is still individually truncated at
+    ``max_character_length``.
+
     Args:
         iteration: The iteration to format
-        max_character_length: The maximum character length of the result
+        max_character_length: Per-block cap on the formatted execution
+            result. Longer outputs are tail-trimmed.
 
     Returns:
-        A list of messages to add to the next prompt
+        A list of messages to add to the next prompt — always length 1
+        (just the assistant) when no code was run, or length 2 (assistant
+        + one combined user reply) otherwise.
     """
     messages = [{"role": "assistant", "content": iteration.response}]
 
-    for code_block in iteration.code_blocks:
-        code = code_block.code
-        result = code_block.result
-        result = format_execution_result(result)
+    parts = []
+    multi = len(iteration.code_blocks) > 1
+    for i, code_block in enumerate(iteration.code_blocks):
+        result = format_execution_result(code_block.result)
         if len(result) > max_character_length:
             result = (
                 result[:max_character_length]
                 + f"... + [{len(result) - max_character_length} chars...]"
             )
+        header = f"REPL output (block {i + 1}):" if multi else "REPL output:"
+        parts.append(f"{header}\n{result}")
 
-        execution_message = {
-            "role": "user",
-            "content": f"Code executed:\n```python\n{code}\n```\n\nREPL output:\n{result}",
-        }
-        messages.append(execution_message)
+    if parts:
+        messages.append({"role": "user", "content": "\n\n".join(parts)})
     return messages