Overview of oolong
mainRLMTrainEnv interface for training purposes.repository·main·Indexed 24 days ago
https://github.com/alexzhang13/rlmA 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.
RLMTrainEnv interface for training purposes.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).An RLM completion consists of three cooperating components:
rlm/core/rlm.py): The main loop that drives the iterative reasoning process.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.rlm/environments/local_repl.py): The Python execution environment where model-generated code runs using exec() in the same process.Execution Flow:
RLM.completion(prompt) spawns an LMHandler (TCP server on a random localhost port) and a LocalREPL.repl code blocks, executes them in LocalREPL, and appends stdout/stderr to the history.answer["ready"] is True or limits are exceeded.RLM operates using a recursive execution model where rlm.completion() calls can trigger child RLM instances.
Execution Flow Summary:
rlm.completion() spawns an LMHandler (a local TCP socket server on an auto-assigned port) and a LocalREPL context.```repl).LocalREPL.execute_code() method runs the generated code using Python's exec() within the same process.rlm_query(), the system checks the current depth. If depth < max_depth, a new child RLM is spawned with its own LMHandler and LocalREPL.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 userRLM 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 setupFor 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 setupRLM supports three primary execution environments:
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",
)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,
},
)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).
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 devTo use Modal Sandboxes as the REPL environment, install the modal library and authenticate your account:
uv add modal
modal setupTo 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("...")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:
uv.# 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 .