Recursive Language Models (RLMs)

repository·main·Indexed 24 days ago

https://github.com/alexzhang13/rlm

A framework enabling language models to handle near-infinite context through programmatic interaction with a REPL environment, supporting recursive sub-calls and code execution. It features a task-agnostic inference paradigm via the RLM class, support for multiple execution environments (local, Docker, Modal, Prime Intellect, Daytona, e2b), and a training harness integrated with prime-rl. Includes tools for trajectory logging and visualization of code and sub-LM calls.

Tokens
13K
Snippets
30
Records
79
Agent score
90%

What's inside rlms

  1. Overview of rlm-train architecture

    main

    The rlm-train package provides a verifiers-compatible training harness for rlm.RLM at depth=1. It integrates with prime-rl for end-to-end RL training. Key components include:

    • src/rlm_train/: Contains the environment, rubric, sub-LM proxy, and subprocess REPL worker.
    • environments/oolong/: An example environment for OOLONG synthetic long-context QA.
    • configs/: Contains example RL configurations (e.g., rlm-qwen3-30b-example.toml).
  2. Understand the RLM Architecture

    main

    An RLM completion consists of three cooperating components:

    1. RLM (rlm/core/rlm.py): The main loop that drives the iterative reasoning process.
    2. LMHandler (rlm/core/lm_handler.py): A per-completion TCP server that routes LM API calls. It allows code in isolated environments (like Docker or Modal) to communicate with the host process via TCP.
    3. LocalREPL (rlm/environments/local_repl.py): The Python execution environment where model-generated code runs using exec() in the same process.

    Execution Flow:

    1. RLM.completion(prompt) spawns an LMHandler (TCP server on a random localhost port) and a LocalREPL.
    2. The RLM iterates: sends history to the LM backend, extracts repl code blocks, executes them in LocalREPL, and appends stdout/stderr to the history.
    3. The loop repeats until answer["ready"] is True or limits are exceeded.
    4. The handler and environment are torn down.
  3. Understand the RLM request flow and execution model

    main

    RLM operates using a recursive execution model where rlm.completion() calls can trigger child RLM instances.

    Execution Flow Summary:

    1. Initialization: rlm.completion() spawns an LMHandler (a local TCP socket server on an auto-assigned port) and a LocalREPL context.
    2. Code Generation: The Language Model (LM) generates code blocks (e.g., using ```repl).
    3. Execution: The LocalREPL.execute_code() method runs the generated code using Python's exec() within the same process.
    4. Recursion: If the code calls rlm_query(), the system checks the current depth. If depth < max_depth, a new child RLM is spawned with its own LMHandler and LocalREPL.
    5. Completion: Once the loop finishes, the LMHandler stops and the RLMChatCompletion is returned to the user.
    User: rlm.completion("Analyze this data")
     │
     ▼
    RLM (depth=0)
     ├─ _spawn_completion_context()
     │   ├─ LMHandler #1 starts on port 52301
     │   └─ LocalREPL #1 created with context="Analyze this data"
     │
     ├─ Iteration 1: LM generates code
     │  ```repl
     │  answer = rlm_query("What patterns exist in: " + context[:5000])
     │  ```
     │
     │  └─ LocalREPL.execute_code() runs the code via exec()
     │      │
     │      ├─ rlm_query() → _rlm_query() → subcall_fn()
     │      │   │
     │      │   └─ RLM._subcall("What patterns exist in: ...")
     │      │       │
     │      │       ├─ depth=1 < max_depth=2, so create child RLM
     │      │       │
     │      │       └─ Child RLM (depth=1)
     │      │           ├─ LMHandler #2 on port 52302
     │      │           ├─ LocalREPL #2 with context="What patterns..."
     │      │           │
     │      │           ├─ Child iteration 1: LM generates code
     │      │           │   │  result = llm_query("Extract key metrics: " + context)
     │      │           │   │
     │      │           │   └─ llm_query() → TCP to Handler #2 → LM API → response
     │      │           │   │
     │      │           │   ├─ Child iteration 2: LM sets answer["content"]=result, answer["ready"]=True
     │      │           │   │
     │      │           │   └─ Returns RLMChatCompletion to parent
     │      │           │
     │      │           └─ child_response = child_completion.response
     │      │       
     │      │ └─ Iteration 2: LM uses child_response, sets answer["content"]=final, answer["ready"]=True
     │      │
     │      └─ LMHandler #1 stops
     │      └─ Returns RLMChatCompletion to user
  4. Install Modal and Docker support for RLM

    main

    RLM supports optional extensions for cloud-based or containerized execution:

    For cloud-based sandboxed execution, install the [modal] extra and authenticate:

    uv pip install -e ".[modal]"
    modal setup

    Docker Support

    For containerized execution, ensure Docker is installed and running on your host machine.

    docker --version
    # Install Modal extra
    uv pip install -e ".[modal]"
    
    # Authenticate Modal
    modal setup
  5. Configure RLM execution environments

    main

    RLM supports three primary execution environments:

    Local (Default)

    Code runs in the same Python process with sandboxed builtins. Fast, but less isolation.

    rlm = RLM(
        backend="openai",
        backend_kwargs={"model_name": "gpt-4o"},
        environment="local",
    )

    Docker

    Code runs in a Docker container with full isolation. A host-side proxy handles LM access so the container doesn't need API keys. Supports persistent=True and compaction=True.

    rlm = RLM(
        backend="openai",
        backend_kwargs={"model_name": "gpt-4o"},
        environment="docker",
        environment_kwargs={
            "image": "python:3.11-slim",
        },
    )

    Code runs in Modal's cloud sandboxes for scalable, fully isolated execution.

    rlm = RLM(
        backend="openai",
        backend_kwargs={"model_name": "gpt-4o"},
        environment="modal",
        environment_kwargs={
            "app_name": "my-rlm-app",
            "timeout": 600,
        },
    )
  6. Train RLMs using the training harness

    main

    The repository includes a training environment based on Prime Intellect's prime-rl. The training logic is located in the training/ folder. It exposes rlm.RLM as a verifiers Environment.

    To add a new training environment, author a verifiers environment that wraps your task and reference it from a configuration file (e.g., a .toml file).

  7. Run the visualizer development server

    main

    To start the RLM trajectory visualizer locally, run the development server using your preferred package manager. Once running, access the interface at http://localhost:3000.

    npm run dev
    # or
    yarn dev
    # or
    pnpm dev
    # or
    bun dev
  8. Enable logging for RLM sessions

    main

    To capture trajectory data, create an instance of RLMLogger and pass it to the RLM constructor. This will save logs in JSON-lines (.jsonl) format to the specified directory. Logs are named using the pattern rlm_TIMESTAMP_UUID.jsonl.

    from rlm import RLM
    from rlm.logger import RLMLogger
    
    # Create logger
    logger = RLMLogger(log_dir="./logs")
    
    rlm = RLM(
        backend="openai",
        backend_kwargs={"model_name": "gpt-4o"},
        logger=logger,
        verbose=True,
    )
    
    result = rlm.completion("...")
  9. Install RLM using uv

    main

    To install RLM, it is recommended to use uv. Ensure you have Python 3.11 or higher installed. Follow these steps to set up your environment:

    1. Install uv.
    2. Initialize a project and create a virtual environment with Python 3.12.
    3. Activate the environment.
    4. Install RLM in editable mode.
    # Install uv
    curl -LsSf https://astral.sh/uv/install.sh | sh
    
    # Create and activate virtual environment
    uv init && uv venv --python 3.12
    source .venv/bin/activate
    
    # Install RLM in editable mode
    uv pip install -e .