rlinf 0.2__py3-none-any.whl

rlinf/__init__.py

@@ -0,0 +1,17 @@
+# Copyright 2025 The RLinf Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .utils.omega_resolver import omegaconf_register
+
+omegaconf_register()

rlinf/agents/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2025 The RLinf Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

rlinf/agents/mas_search/__init__.py

@@ -0,0 +1,18 @@
+# Copyright 2026 The RLinf Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from rlinf.agents.mas_search.mas_search_agent_loop import MasSearchAgentLoopWorker
+from rlinf.agents.searchr1.search_tool_worker import SearchToolWorker
+
+__all__ = ["MasSearchAgentLoopWorker", "SearchToolWorker"]

rlinf/agents/mas_search/mas_search_agent_loop.py

@@ -0,0 +1,222 @@
+# Copyright 2026 The RLinf Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import asyncio
+import copy
+import json
+import re
+from uuid import uuid4
+
+from omegaconf import DictConfig
+
+from rlinf.data.tool_call.tool_io_struct import (
+    ToolChannelRequest,
+    ToolChannelResponse,
+    ToolRequest,
+    ToolResponse,
+)
+from rlinf.scheduler import Channel
+from rlinf.utils.placement import ModelParallelComponentPlacement
+from rlinf.workers.agent.agent_loop import (
+    AgentLoopOutput,
+    MultiAgentLoopOutput,
+    MultiAgentLoopWorker,
+)
+
+
+class MasSearchAgentLoopWorker(MultiAgentLoopWorker):
+    """
+    Agent loop worker that combines search-r1's <search>keyword</search> extraction
+    logic with multi agent system's component structure.
+    """
+
+    def __init__(
+        self,
+        cfg: DictConfig,
+        placement: ModelParallelComponentPlacement,
+    ):
+        super().__init__(cfg, placement)
+        self.max_prompt_len = int(self.cfg.data.max_prompt_length)
+        max_total_len = int(self.cfg.actor.model.encoder_seq_length)
+        self.max_resp_len = max(1, max_total_len - self.max_prompt_len)
+
+        # Search tool call tokens
+        self.tool_call_start_token: str = "<search>"
+        self.tool_call_end_token: str = "</search>"
+        self.tool_call_regex = re.compile(r"<search>(.*?)</search>", re.DOTALL)
+
+        # Max turns for multi-turn interaction
+        self.max_turns = self.cfg.agentloop.get("max_turns", 5)
+        self.return_logprobs = not cfg.algorithm.recompute_logprobs
+
+    async def state_less_tool_call_with_channel(
+        self,
+        input_channel: Channel,
+        output_channel: Channel,
+        tool_name: str,
+        tool_args: dict,
+    ) -> ToolChannelResponse:
+        """state-less tool call with channel, used for demo"""
+        session_id = uuid4().hex
+        await input_channel.put(
+            ToolChannelRequest(
+                session_id=session_id,
+                request_type="execute",
+                tool_name=tool_name,
+                tool_args=tool_args,
+            ),
+            async_op=True,
+        ).async_wait()
+        return await output_channel.get(session_id, async_op=True).async_wait()
+
+    async def tool_call(self, tool_request: ToolRequest) -> ToolResponse:
+        tool_name, tool_args = tool_request.name, tool_request.arguments
+        tool_channel_info = self.tool_channel_info_map[self.tool_name_map[tool_name]]
+        channel_response = await self.state_less_tool_call_with_channel(
+            tool_channel_info.input_channel,
+            self.tool_worker_output_channel,
+            tool_name,
+            tool_args,
+        )
+
+        # no failure in this demo
+        assert channel_response.success
+        if isinstance(channel_response.result, (list, dict)):
+            result_text = json.dumps(channel_response.result)
+        else:
+            result_text = str(channel_response.result)
+        return ToolResponse(
+            text=result_text,
+        )
+
+    async def extract_tool_calls(self, response_text) -> tuple[str, list[ToolRequest]]:
+        """
+        Extract tool calls from response text using <search>keyword</search> format.
+        """
+        if (
+            self.tool_call_start_token not in response_text
+            or self.tool_call_end_token not in response_text
+        ):
+            return response_text, []
+        matches = self.tool_call_regex.findall(response_text)
+        function_calls = []
+        if matches:
+            match = matches[-1].strip()
+            function_calls.append(
+                ToolRequest(name="search", arguments={"keyword": match})
+            )
+
+        # remaining text exclude tool call tokens
+        content = self.tool_call_regex.sub("", response_text)
+
+        return content, function_calls
+
+    async def run_one_query(self, prompt_ids: list[int], *, answer) -> AgentLoopOutput:
+        output_buffer = []
+        orig_prompt_ids = copy.deepcopy(prompt_ids)
+        trace_prints = []
+        response_mask = []
+        all_response_ids = []
+        for _ in range(self.max_turns):
+            # Generate response from LLM
+            max_resp_len = self.max_resp_len - (len(prompt_ids) - len(orig_prompt_ids))
+
+            generate_result = await self.generate(
+                prompt_ids, sampling_params={"max_new_tokens": max_resp_len}
+            )
+            generate_prompt_ids = copy.deepcopy(prompt_ids)
+            response_ids = generate_result["output_ids"]
+            if len(response_ids) > max_resp_len:
+                response_ids = response_ids[:max_resp_len]
+            response_text = self.tokenizer.decode(response_ids)
+
+            # # split </search> manually
+            # if "</search>" in response_text:
+            #     response_text = response_text.split("</search>")[0] + "</search>"
+            #     response_ids = self.tokenizer.encode(response_text)
+
+            output_buffer.append(
+                AgentLoopOutput(
+                    prompt_ids=copy.deepcopy(prompt_ids),
+                    response_ids=copy.deepcopy(response_ids),
+                    prompt_text=self.tokenizer.decode(prompt_ids),
+                    response_text=response_text,
+                    response_mask=response_mask,
+                    response_logprobs=generate_result["logprobs"]
+                    if self.return_logprobs
+                    else None,
+                )
+            )
+
+            prompt_ids += response_ids
+            all_response_ids.extend(response_ids)
+            response_mask += [1] * len(response_ids)
+
+            if len(response_ids) == max_resp_len:
+                break
+
+            # Extract tool calls from response text
+            content, function_calls = await self.extract_tool_calls(response_text)
+
+            if function_calls == []:
+                break
+
+            # Execute tools in parallel with history propagation
+            tasks = []
+            for tool_request in function_calls:
+                tasks.append(self.tool_call(tool_request))
+            tool_responses: list[ToolResponse] = await asyncio.gather(*tasks)
+
+            # Convert tool responses to messages and tokenize
+            tool_messages = []
+            for tool_response in tool_responses:
+                message = {"role": "tool", "content": tool_response.text}
+                tool_messages.append(message)
+            tool_response_ids = self.tokenizer.encode(
+                tool_messages[0]["content"], add_special_tokens=False
+            )
+            max_tool_resp_len = self.max_resp_len - (
+                len(prompt_ids) - len(orig_prompt_ids)
+            )
+            if len(tool_response_ids) > max_tool_resp_len:
+                # task_failed = True
+                break
+            prompt_ids += tool_response_ids
+            all_response_ids.extend(tool_response_ids)
+            response_mask += [0] * len(tool_response_ids)
+
+            if self.print_outputs:
+                # add anything you want to print
+                trace_prints.append(
+                    {
+                        "decode_prompt": self.tokenizer.decode(generate_prompt_ids),
+                        "generate": response_text,
+                        "tool_resp": tool_messages,
+                    }
+                )
+        # Build complete response text from all turns
+        complete_response_text = self.tokenizer.decode(all_response_ids)
+
+        # Extract final answer from complete response using searchr1's extract_solution
+        from rlinf.algorithms.rewards.searchr1 import compute_score
+
+        reward_score = compute_score(complete_response_text, answer)
+
+        for single_turn_output in output_buffer:
+            single_turn_output.reward_score = reward_score
+
+        return MultiAgentLoopOutput(
+            single_turn_outputs=output_buffer,
+            trace_prints=[],
+        )

rlinf/agents/rstar2/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2026 The RLinf Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

rlinf/agents/rstar2/http_code_judge_tool.py

@@ -0,0 +1,331 @@
+# Copyright 2026 The RLinf Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import base64
+import time
+from typing import Callable, Optional
+
+from rlinf.data.tool_call.tool_io_struct import ToolChannelRequest, ToolChannelResponse
+
+
+class ToolBase:
+    name = None
+
+    def __init__(self, cfg):
+        self.cfg = cfg
+
+    async def execute(
+        self, request: ToolChannelRequest, **kwargs
+    ) -> ToolChannelResponse:
+        """Execute the tool call."""
+        raise NotImplementedError()
+
+    def tool_schema(self) -> dict:
+        """
+        A JSON Schema, giving the name, description and argument types for the tool.
+        Ref: https://huggingface.co/docs/transformers/en/chat_extras#json-schemas
+        """
+        raise NotImplementedError()
+
+    def validate(self, request: ToolChannelRequest) -> Optional[ToolChannelResponse]:
+        """Validate the request call the right tool and the schema is right."""
+        raise NotImplementedError()
+
+
+class CodeJudgeToolBase(ToolBase):
+    name = None
+
+    def __init__(self, cfg):
+        super().__init__(cfg=cfg)
+        self.url = f"http://{self.cfg.tools.codejudge.host_addr}:{self.cfg.tools.codejudge.host_port}/run/long-batch"
+
+    def _postprocess(self, result):
+        if result["run_success"] and result["success"]:
+            output_parts = []
+            output_parts.append("Tool call success")
+            if result["stdout"]:
+                output_parts.append(f"stdout: {result['stdout']}")
+            if result["stderr"]:
+                output_parts.append(f"stderr: {result['stderr']}")
+            output_parts.append(f"execution time: {result['cost']:.2f}s")
+            result = "\n".join(output_parts)
+            return ToolChannelResponse(success=True, result=result)
+        else:
+            output_parts = []
+            output_parts.append("Tool call failure")
+            output_parts.append(f"reason: {result['reason']}")
+            if result["stdout"]:
+                output_parts.append(f"stdout: {result['stdout']}")
+            if result["stderr"]:
+                output_parts.append(f"stderr: {result['stderr']}")
+            output_parts.append(f"execution time: {result['cost']:.2f}s")
+            result = "\n".join(output_parts)
+            return ToolChannelResponse(success=False, result=result)
+
+
+code_template_setup = '''
+import os
+import base64
+import sys
+import ast
+import traceback
+from typing import Optional, Any
+import linecache
+from types import CodeType
+from contextlib import redirect_stdout, redirect_stderr
+from io import StringIO
+
+class CodeExecutionError(Exception):
+    """Custom exception for code execution errors with line information"""
+    def __init__(self, original_error: Exception, code: str, line_offset: int = 0):
+        self.original_error = original_error
+        self.code = code
+        self.line_offset = line_offset
+
+        # Get error line number
+        if hasattr(original_error, 'lineno'):
+            self.lineno = original_error.lineno
+        else:
+            tb = getattr(original_error, '__traceback__', None)
+            if tb:
+                while tb.tb_next:
+                    tb = tb.tb_next
+                self.lineno = tb.tb_lineno
+            else:
+                self.lineno = -1
+
+        # Adjust line number for code segment
+        if self.lineno != -1:
+            self.lineno += line_offset
+
+        # Format error message
+        error_type = type(original_error).__name__
+        error_msg = str(original_error)
+
+        if self.lineno != -1:
+            # Get the problematic line
+            lines = code.splitlines()
+            if 0 <= self.lineno - 1 < len(lines):
+                error_line = lines[self.lineno - 1]
+                # Create error message with line information
+                super().__init__(f"{error_type} at line {self.lineno}: {error_msg}\\n  {error_line}")
+                return
+
+        super().__init__(f"{error_type}: {error_msg}")
+
+class PersistentExecutor:
+    def __init__(self):
+        self.exec_globals = {
+            '__name__': '__main__',
+            '__file__': '<string>',
+            '__builtins__': __builtins__
+        }
+
+    def split_code(self, code: str) -> tuple[str, Optional[str]]:
+        """
+        Intelligently split code into main body and last expression
+
+        Args:
+            code: The source code string
+
+        Returns:
+            tuple[str, Optional[str]]: (main code body, last expression if exists)
+        """
+        try:
+            # Parse code into AST
+            tree = ast.parse(code)
+            if not tree.body:
+                return code, None
+
+            # Check if the last node is a pure expression (not a call)
+            last_node = tree.body[-1]
+            if isinstance(last_node, ast.Expr):
+                # Get the line range of the last expression
+                last_expr_start = last_node.lineno
+                last_expr_end = last_node.end_lineno if hasattr(last_node, 'end_lineno') else last_node.lineno
+
+                # Split the code
+                lines = code.splitlines()
+                main_code = '\\n'.join(lines[:last_expr_start-1])
+                last_expr = '\\n'.join(lines[last_expr_start-1:last_expr_end])
+                return main_code, last_expr
+        except SyntaxError as e:
+            raise CodeExecutionError(e, code)
+        return code, None
+
+    def execute_code(self, code: str, replay_history_code: bool) -> None:
+        """
+        Execute code while maintaining persistent environment state.
+        If the last line is an expression, its value will be printed to stdout.
+
+        Args:
+            code: The source code string to execute
+            replay_history_code: If True, suppress stdout and stderr output
+        """
+        try:
+            # Split code intelligently
+            main_code, last_expr = self.split_code(code)
+
+            # Set up output redirection if replay_history_code is True
+            if replay_history_code:
+                stdout_capture = StringIO()
+                stderr_capture = StringIO()
+                stdout_context = redirect_stdout(stdout_capture)
+                stderr_context = redirect_stderr(stderr_capture)
+            else:
+                stdout_context = redirect_stdout(sys.stdout)
+                stderr_context = redirect_stderr(sys.stderr)
+
+            # Execute main code body
+            if main_code:
+                try:
+                    # Compile code to get better error line numbers
+                    compiled_code = compile(main_code, '<string>', 'exec')
+                    with stdout_context, stderr_context:
+                        exec(compiled_code, self.exec_globals)
+                except Exception as e:
+                    raise CodeExecutionError(e, main_code)
+
+            # If there's a last expression, try to evaluate and print it
+            if last_expr:
+                try:
+                    # Compile expression to get better error line numbers
+                    compiled_expr = compile(last_expr, '<string>', 'eval')
+                    with stdout_context, stderr_context:
+                        last_value = eval(compiled_expr, self.exec_globals)
+
+                    # Only print the result if not in replay mode
+                    if last_value is not None and not replay_history_code:
+                        print(repr(last_value), file=sys.stdout)
+                except Exception as e:
+                    # Try executing as statement if evaluation fails
+                    try:
+                        compiled_stmt = compile(last_expr, '<string>', 'exec')
+                        with stdout_context, stderr_context:
+                            exec(compiled_stmt, self.exec_globals)
+                    except Exception as e:
+                        # Calculate line offset for the last expression
+                        line_offset = len(main_code.splitlines()) if main_code else 0
+                        raise CodeExecutionError(e, last_expr, line_offset)
+
+        except Exception as e:
+            if replay_history_code:
+                return
+            if isinstance(e, CodeExecutionError):
+                print(str(e), file=sys.stderr)
+            else:
+                traceback.print_exc(file=sys.stderr)
+            os._exit(1)
+            return
+
+persistent_executor = PersistentExecutor()
+'''
+
+
+code_template_exec = """
+code_to_execute = base64.b64decode("{}".encode()).decode()
+persistent_executor.execute_code(code_to_execute, replay_history_code={})
+"""
+
+
+class PythonTool(CodeJudgeToolBase):
+    name = "python_code_with_standard_io"
+
+    def __init__(self, cfg):
+        super().__init__(cfg=cfg)
+
+    async def execute(
+        self,
+        request: ToolChannelRequest,
+        send_request_func: Callable[[str, dict], ToolChannelResponse],
+    ):
+        err_msg = self.validate(request)
+        if err_msg:
+            return err_msg
+
+        # convert the code to the code exec on code-judge
+        code_to_execute = base64.b64encode(
+            request.tool_args.get("code", "").encode()
+        ).decode()
+        final_code = code_template_setup
+        # TODO: add history code here
+        final_code += code_template_exec.format(code_to_execute, "False")
+
+        submission = {
+            "type": "python",
+            "solution": final_code,
+            "input": request.tool_args.get("input", ""),
+        }
+
+        data = {"type": "batch", "submissions": [submission]}
+
+        for retry_time in range(4):
+            try:
+                results = (await send_request_func(self.url, data))["results"]
+                break
+            except Exception as e:
+                print(f"Tool retry time {retry_time}, exception: {e}")
+                time.sleep(1)
+        else:
+            raise RuntimeError("Tool call failed after retries")
+        assert len(results) == 1, f"{results}"
+        return self._postprocess(results[0])
+
+    def tool_schema(self) -> dict:
+        return {
+            "type": "function",
+            "function": {
+                "name": "python_code_with_standard_io",
+                "description": "Execute Python code with standard input and capture standard output. This function takes a Python code string and an input string, provides the input string through standard input (stdin) to the code, and captures and returns any output produced through standard output (stdout). If the executed code raises an exception, the error message will be captured and returned instead.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "code": {
+                            "type": "string",
+                            "description": "A string containing Python code to be executed. The code can read from standard input using the input() function.",
+                        },
+                        "input": {
+                            "type": "string",
+                            "description": "A string that will be provided as standard input to the code when it calls input().",
+                        },
+                    },
+                    "required": ["code", "input"],
+                },
+            },
+        }
+
+    def validate(self, request) -> Optional[ToolChannelResponse]:
+        tool_args = request.tool_args
+
+        assert request.tool_name == self.name, (
+            f"Name mismatch, {self.name} != {request.tool_name}"
+        )
+
+        if not isinstance(tool_args, dict):
+            return ToolChannelResponse(
+                success=False,
+                result="Error when executing tool: run_tool_calls_on_server_async failed for1 tool calls after 4 attempts.",
+            )
+
+        if "code" in tool_args and not isinstance(tool_args["code"], str):
+            return ToolChannelResponse(
+                success=False,
+                result="Error when executing tool: run_tool_calls_on_server_async failed for1 tool calls after 4 attempts.",
+            )
+
+        if "input" in tool_args and not isinstance(tool_args["input"], str):
+            return ToolChannelResponse(
+                success=False,
+                result="Error when executing tool: run_tool_calls_on_server_async failed for1 tool calls after 4 attempts.",
+            )