From 38917fb4f3c0ada3b4edf621df942ffe607aaa27 Mon Sep 17 00:00:00 2001 From: z0o0ey Date: Sat, 1 Aug 2026 15:45:44 +0800 Subject: [PATCH 1/8] wip --- .gitignore | 2 + cookbook/rl/agentenv/Dockerfile | 26 + cookbook/rl/agentenv/agentenv_grpo.py | 370 ++++++++++++ cookbook/rl/agentenv/agentenv_grpo.sh | 54 ++ cookbook/rl/agentenv/tools.py | 137 +++++ cookbook/rl/code_rl/Dockerfile | 26 + cookbook/rl/code_rl/README.md | 264 +++++++++ cookbook/rl/code_rl/backends/__init__.py | 51 ++ cookbook/rl/code_rl/backends/agentenv.py | 156 +++++ cookbook/rl/code_rl/backends/openenv.py | 153 +++++ cookbook/rl/code_rl/common_args.sh | 36 ++ cookbook/rl/code_rl/run_agentenv.sh | 45 ++ cookbook/rl/code_rl/run_openenv.sh | 34 ++ cookbook/rl/code_rl/serve.sh | 52 ++ cookbook/rl/code_rl/server_app.py | 113 ++++ cookbook/rl/code_rl/train.py | 408 +++++++++++++ cookbook/rl/deploy/README.md | 48 ++ cookbook/rl/deploy/openenv-server.service | 44 ++ cookbook/rl/openenv_code/openenv_code_grpo.py | 365 ++++++++++++ cookbook/rl/openenv_code/openenv_code_grpo.sh | 48 ++ cookbook/rl/openenv_code/serve.sh | 52 ++ cookbook/rl/openenv_code/server_app.py | 113 ++++ cookbook/rl/openenv_code/tools.py | 138 +++++ docs/source_en/Components/Agentic/Envs.md | 120 +++- .../Agentic/Multi-Turn-Tool-Usage.md | 17 +- .../Usage Guide/Agentic-RL-Best-Practices.md | 305 ++++++++++ .../Usage Guide/Agentic-RL-Deployment.md | 176 ++++++ .../Usage Guide/Agentic-RL-Sandbox.md | 281 +++++++++ docs/source_en/index.rst | 3 + docs/source_zh/index.rst | 1 + ...62\344\270\216\350\256\255\347\273\203.md" | 547 ++++++++++++++++++ .../\347\273\204\344\273\266/Agentic/Envs.md" | 120 +++- .../Agentic/Multi-Turn-Tool-Usage.md" | 17 +- src/twinkle_agentic/envs/__init__.py | 3 +- src/twinkle_agentic/envs/agentenv.py | 363 ++++++++++++ src/twinkle_agentic/envs/openenv.py | 322 ++++++++++- 36 files changed, 4960 insertions(+), 50 deletions(-) create mode 100644 cookbook/rl/agentenv/Dockerfile create mode 100644 cookbook/rl/agentenv/agentenv_grpo.py create mode 100644 cookbook/rl/agentenv/agentenv_grpo.sh create mode 100644 cookbook/rl/agentenv/tools.py create mode 100644 cookbook/rl/code_rl/Dockerfile create mode 100644 cookbook/rl/code_rl/README.md create mode 100644 cookbook/rl/code_rl/backends/__init__.py create mode 100644 cookbook/rl/code_rl/backends/agentenv.py create mode 100644 cookbook/rl/code_rl/backends/openenv.py create mode 100644 cookbook/rl/code_rl/common_args.sh create mode 100644 cookbook/rl/code_rl/run_agentenv.sh create mode 100644 cookbook/rl/code_rl/run_openenv.sh create mode 100644 cookbook/rl/code_rl/serve.sh create mode 100644 cookbook/rl/code_rl/server_app.py create mode 100644 cookbook/rl/code_rl/train.py create mode 100644 cookbook/rl/deploy/README.md create mode 100644 cookbook/rl/deploy/openenv-server.service create mode 100644 cookbook/rl/openenv_code/openenv_code_grpo.py create mode 100644 cookbook/rl/openenv_code/openenv_code_grpo.sh create mode 100644 cookbook/rl/openenv_code/serve.sh create mode 100644 cookbook/rl/openenv_code/server_app.py create mode 100644 cookbook/rl/openenv_code/tools.py create mode 100644 docs/source_en/Usage Guide/Agentic-RL-Best-Practices.md create mode 100644 docs/source_en/Usage Guide/Agentic-RL-Deployment.md create mode 100644 docs/source_en/Usage Guide/Agentic-RL-Sandbox.md create mode 100644 "docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/Agentic RL\351\203\250\347\275\262\344\270\216\350\256\255\347\273\203.md" create mode 100644 src/twinkle_agentic/envs/agentenv.py diff --git a/.gitignore b/.gitignore index 222caa50b..c643ab30e 100644 --- a/.gitignore +++ b/.gitignore @@ -39,6 +39,8 @@ MANIFEST # before PyInstaller builds the exe, so as to inject date/other infos into it. *.manifest *.spec +.cursor +.qoder # Installer logs pip-log.txt diff --git a/cookbook/rl/agentenv/Dockerfile b/cookbook/rl/agentenv/Dockerfile new file mode 100644 index 000000000..f4ab87388 --- /dev/null +++ b/cookbook/rl/agentenv/Dockerfile @@ -0,0 +1,26 @@ +# Sandbox template for cookbook/rl/agentenv (code-writing RL on MBPP). +# +# A general-purpose Python code interpreter: the model writes functions, runs +# them, and the training loop replays hidden unit tests inside this sandbox. +# Unlike OpenEnv's AST-interpreter backend, this is a real CPython process in a +# Firecracker microVM, so the full language (assert / try / subprocess / files) +# and arbitrary pip packages are available. +# +# Build this template ONCE against your AgentENV deployment (offline step, +# requires the aenv CLI authenticated via `aenv auth`): +# +# aenv build cookbook/rl/agentenv/Dockerfile -t twinkle-code \ +# --cpu-count 1 --memory-mb 1024 +# aenv template watch # wait until the build is ready +# +# aenv build understands FROM / RUN / ENV / WORKDIR / USER instructions and +# bakes the result into a snapshot, so sandboxes boot in ~50ms with all +# dependencies pre-installed. Rebuild only when this file changes. +FROM python:3.11-slim + +# Use a pip mirror if your sandbox network requires it, e.g.: +# ENV PIP_INDEX_URL=https://mirrors.aliyun.com/pypi/simple +RUN pip install --no-cache-dir numpy sympy + +ENV PYTHONUNBUFFERED=1 +WORKDIR /workspace diff --git a/cookbook/rl/agentenv/agentenv_grpo.py b/cookbook/rl/agentenv/agentenv_grpo.py new file mode 100644 index 000000000..adcb36305 --- /dev/null +++ b/cookbook/rl/agentenv/agentenv_grpo.py @@ -0,0 +1,370 @@ +"""Multi-turn GRPO on MBPP with AgentENV Firecracker sandboxes. + +Sandbox counterpart of ``cookbook/rl/openenv_code/openenv_code_grpo.py``: the +same code-writing task, but every trajectory gets its own **Firecracker +microVM** instead of a session on a shared OpenEnv server. Compare the two to +see what the execution backend does and does not change. + +Key architectural difference vs the embedded OpenEnv example +(``cookbook/rl/multi_turn/multi_turn_grpo.py``): + - No EnvPool / @remote_class. AgentEnv is a stateless HTTP client; sandbox + placement, load balancing, pause/resume and failover are handled + server-side by AgentENV's gateway/scheduler. The driver just creates one + AgentEnv per trajectory (parallelized with a thread pool, since each + reset() is a blocking HTTP call that boots a sandbox). + +Per-trajectory flow: + 1. reset() boots a sandbox from the pre-built template (~50ms from snapshot). + 2. MultiTurnRollout drives tool calls: run_python executes a self-contained + snippet in the sandbox; submit_solution records the final source. + 3. Reward = MBPP unit-test pass rate, measured by replaying the hidden tests + inside the same sandbox after the rollout (see ``tools.run_tests``). + 4. GRPO advantages are group-relative across NUM_GENERATIONS rollouts of + the same problem. + +Prerequisites (offline, once): + 1. Deploy AgentENV (single server or gateway+scheduler cluster). + 2. Build the sandbox template from the Dockerfile in this folder: + aenv build cookbook/rl/agentenv/Dockerfile -t twinkle-code + 3. pip install e2b (on the training side) + +Usage: + AENV_API_URL=http://:8000 sh agentenv_grpo.sh +""" +import os +from concurrent.futures import ThreadPoolExecutor +from typing import Any, Dict, List, Tuple + +from peft import LoraConfig + +import twinkle +from twinkle import DeviceGroup, DeviceMesh, get_device_placement, get_logger +from twinkle.advantage import GRPOAdvantage +from twinkle.checkpoint_engine import CheckpointEngineManager +from twinkle.cli import CLI +from twinkle.data_format import SamplingParams +from twinkle.metric import CompletionRewardMetric +from twinkle.model import TransformersModel +from twinkle.processor import InputProcessor +from twinkle.sampler import vLLMSampler +from twinkle.template import Qwen3_5Template +from twinkle_agentic.envs import AgentEnv, EnvTool +from twinkle_agentic.rollout.multi_turn import MultiTurnRollout +from twinkle_agentic.tools.tool_manager import ToolManager +from tools import SYSTEM_PROMPT, TOOL_SCHEMA, register_tools, run_tests + +logger = get_logger() +args = CLI.from_args() + +# ========== Configuration ========== +MODEL_ID = args.model.model_id or 'ms://Qwen/Qwen3.5-4B' + +MODEL_GPUS = args.infra.model_gpus or 4 +SAMPLER_GPUS = args.infra.sampler_gpus or 4 +NUM_GPUS = MODEL_GPUS + SAMPLER_GPUS + +NUM_GENERATIONS = args.rl.num_generations or 8 +MAX_NEW_TOKENS = args.sampling.max_tokens or 2048 +LEARNING_RATE = args.optimizer.learning_rate or 1e-5 +MAX_STEPS = args.training.max_steps or 1000 +BATCH_SIZE = args.training.batch_size or 4 +MINI_BATCH_SIZE = args.training.mini_batch_size or 8 +MICRO_BATCH_SIZE = args.training.micro_batch_size or 2 +GRADIENT_ACCUMULATION_STEPS = args.training.gradient_accumulation_steps or 1 +ADAPTER_NAME = args.lora.adapter_name or 'default' +SAVE_STEPS = args.training.save_steps or 500 +LORA_RANK = args.lora.lora_r or 16 +MAX_TURNS = int(os.environ.get('MAX_TURNS', '6')) + +# AgentENV deployment. The client only ever needs this ONE address; with a +# multi-node cluster point it at the gateway and node routing is automatic. +AENV_API_URL = os.environ.get('AENV_API_URL', 'http://127.0.0.1:8000') +AENV_TEMPLATE = os.environ.get('AENV_TEMPLATE', 'twinkle-code') +# Sandbox idle timeout: must exceed the slowest full rollout of a batch, +# otherwise mid-episode sandboxes get paused (auto-resume adds latency). +SANDBOX_TIMEOUT = int(os.environ.get('SANDBOX_TIMEOUT', '600')) +# Parallelism for sandbox create/kill HTTP calls in the driver. +ENV_CONCURRENCY = int(os.environ.get('ENV_CONCURRENCY', '16')) + + +# ========== Dataset (MBPP) ========== +def load_mbpp() -> List[Dict[str, Any]]: + """Load MBPP as [{'prompt', 'test_list', 'test_setup_code'}].""" + from modelscope.msdatasets import MsDataset + raw = MsDataset.load('opencompass/mbpp', subset_name='full', split='train') + samples = [] + for row in raw: + test_list = list(row['test_list'] or []) + if not test_list: + continue + # The first assertion is handed to the model as a signature hint: MBPP + # descriptions alone do not pin the function name, and the hidden tests + # call it by name. + prompt = (f"{row['text']}\n\n" + f'Your function must satisfy this call:\n{test_list[0]}') + samples.append({ + 'prompt': prompt, + 'test_list': test_list, + 'test_setup_code': row.get('test_setup_code') or '', + }) + logger.info(f'MBPP loaded: {len(samples)} samples') + return samples + + +# ========== Environment Setup ========== +def make_env() -> AgentEnv: + """One AgentEnv per trajectory; only the two task tools are exposed.""" + env = AgentEnv( + template=AENV_TEMPLATE, + api_url=AENV_API_URL, + sandbox_timeout=SANDBOX_TIMEOUT, + command_timeout=60, + include_default_tools=False, + ) + return register_tools(env) + + +def prepare_trajectories( + samples: List[Dict[str, Any]], + pool: ThreadPoolExecutor, +) -> Tuple[List[Dict[str, Any]], List[ToolManager], List[AgentEnv]]: + """Boot one sandbox per trajectory (in parallel) and build trajectories. + + Mirrors ``prepare_trajectories`` in the OpenEnv examples, except env + creation is a remote HTTP call, so reset() runs on a thread pool. + """ + envs = [make_env() for _ in samples] + # reset() blocks on sandbox boot (~50ms server-side + HTTP RTT); do them + # concurrently so a batch of 32 does not serialize. + list(pool.map(lambda env: env.reset(), envs)) + + trajectories = [] + tool_managers = [] + for sample, env in zip(samples, envs): + env_tools = EnvTool.from_env(env) + tool_managers.append(ToolManager(env_tools)) + trajectories.append({ + 'messages': [ + {'role': 'system', 'content': SYSTEM_PROMPT}, + {'role': 'user', 'content': sample['prompt']}, + ], + 'tools': TOOL_SCHEMA, + }) + return trajectories, tool_managers, envs + + +def close_envs(envs: List[AgentEnv], pool: ThreadPoolExecutor) -> None: + """Kill all sandboxes; best-effort (AgentENV auto-evicts on timeout).""" + list(pool.map(lambda env: env.close(), envs)) + + +def extract_rewards( + envs: List[AgentEnv], + samples: List[Dict[str, Any]], + pool: ThreadPoolExecutor, +) -> Tuple[List[float], List[float]]: + """Score submitted solutions against the hidden MBPP tests. + + Shaped so that a full pass dominates, while a submission that passes some + tests still beats one that passes none — which beats never submitting. + + Returns: + ``(rewards, pass_rates)``. + """ + + def score(pair) -> Tuple[float, float]: + env, sample = pair + passed, total = run_tests(env, sample['test_list'], sample['test_setup_code']) + rate = passed / total if total else 0.0 + if total and passed == total: + return 1.0, rate + if getattr(env, 'submitted_code', None): + return 0.1 + 0.4 * rate, rate + return 0.0, rate + + scored = list(pool.map(score, zip(envs, samples))) + return [s[0] for s in scored], [s[1] for s in scored] + + +# ========== Main ========== +def main(): + device_groups = [ + DeviceGroup(name='model', ranks=list(range(MODEL_GPUS)), device_type='GPU'), + DeviceGroup(name='sampler', ranks=list(range(MODEL_GPUS, NUM_GPUS)), device_type='GPU'), + ] + model_mesh = DeviceMesh.from_sizes(world_size=MODEL_GPUS, dp_size=MODEL_GPUS) + sampler_mesh = DeviceMesh.from_sizes(world_size=SAMPLER_GPUS, dp_size=SAMPLER_GPUS) + twinkle.initialize(mode='ray', nproc_per_node=NUM_GPUS, groups=device_groups, lazy_collect=False) + + lora_config = LoraConfig( + target_modules='all-linear', + r=LORA_RANK, + lora_alpha=LORA_RANK * 2, + lora_dropout=0.05, + ) + + model = TransformersModel( + model_id=MODEL_ID, + device_mesh=model_mesh, + remote_group='model', + ) + model.add_adapter_to_model(ADAPTER_NAME, lora_config, gradient_accumulation_steps=GRADIENT_ACCUMULATION_STEPS) + model.set_optimizer('AdamW', lr=LEARNING_RATE) + model.set_lr_scheduler('CosineAnnealingLR', T_max=MAX_STEPS, eta_min=0) + model.set_loss('GRPOLoss', epsilon=0.2) + model.set_processor(InputProcessor, padding_free=True) + model.set_template('Qwen3_5Template', model_id=MODEL_ID, enable_thinking=False) + + sampler = vLLMSampler( + model_id=MODEL_ID, + engine_args={ + 'gpu_memory_utilization': 0.8, + 'max_model_len': 8192, + 'max_lora_rank': 32, + 'enable_lora': True, + 'enable_tower_connector_lora': True, + }, + device_mesh=sampler_mesh, + remote_group='sampler', + ) + sampler.set_template('Qwen3_5Template', model_id=MODEL_ID, enable_thinking=False) + + rollout_template = Qwen3_5Template(MODEL_ID, max_length=8192, enable_thinking=False) + rollout_template.truncation_strategy = 'delete' + + ckpt_manager = CheckpointEngineManager(model=model, sampler=sampler) + + sampling_params = SamplingParams( + max_tokens=MAX_NEW_TOKENS, num_samples=1, logprobs=1, + temperature=1.0, top_p=0.95, + ) + rollout = MultiTurnRollout( + sampler=sampler, + template=rollout_template, + sampling_params=sampling_params, + max_turns=MAX_TURNS, + ) + + advantage_fn = GRPOAdvantage() + metrics = CompletionRewardMetric() + dataset = load_mbpp() + env_pool = ThreadPoolExecutor(max_workers=ENV_CONCURRENCY) + + optim_step = 0 + sample_cursor = 0 + logger.info(f'Starting AgentENV GRPO: api={AENV_API_URL}, template={AENV_TEMPLATE}') + logger.info(get_device_placement()) + + while optim_step < MAX_STEPS: + metrics.reset() + + # BATCH_SIZE problems x NUM_GENERATIONS contiguous copies (GRPO groups). + batch = [dataset[(sample_cursor + i) % len(dataset)] for i in range(BATCH_SIZE)] + sample_cursor += BATCH_SIZE + expanded = [s for s in batch for _ in range(NUM_GENERATIONS)] + n_traj = len(expanded) + + # 1. Boot sandboxes and build initial trajectories + logger.info(f'[Step {optim_step}] Booting {n_traj} sandboxes...') + expand_prompts, tool_managers, envs = prepare_trajectories(expanded, env_pool) + + try: + # 2. Sync model weights to sampler + ckpt_manager.sync_weights(merge_and_sync=False) + sampler.reset_prefix_cache() + + # 3. Multi-turn rollout with per-trajectory ToolManagers + all_trajectories: List[Dict[str, Any]] = rollout( + expand_prompts, + tool_manager=tool_managers, + ) + + # 4. Score solutions by replaying the hidden tests in each sandbox + total_rewards, pass_rates = extract_rewards(envs, expanded, env_pool) + finally: + # Sandboxes are single-episode; always release them. + close_envs(envs, env_pool) + + all_old_logps: List[List[float]] = [] + all_completion_lengths: List[int] = [] + n_turns_per_rollout: List[int] = [] + for traj in all_trajectories: + logprobs = traj.get('logprobs') or [] + all_old_logps.append([lp[0][1] for lp in logprobs] if logprobs else []) + labels = traj.get('labels') or [] + all_completion_lengths.append(sum(1 for l in labels if l != -100)) + n_turns_per_rollout.append(int(traj.get('turns') or 0)) + + # 5. Group-relative advantages + advantages = advantage_fn( + total_rewards, num_generations=NUM_GENERATIONS, scale='group', + ).tolist() + + # 6. Log metrics + metrics.accumulate( + completion_lengths=all_completion_lengths, + rewards={'total': total_rewards}, + ) + avg_reward = sum(total_rewards) / len(total_rewards) if total_rewards else 0.0 + solve_rate = sum(1 for r in total_rewards if r >= 1.0) / max(len(total_rewards), 1) + avg_pass_rate = sum(pass_rates) / len(pass_rates) if pass_rates else 0.0 + avg_turns = sum(n_turns_per_rollout) / len(n_turns_per_rollout) if n_turns_per_rollout else 0.0 + logger.info(f'[Step {optim_step}] avg_reward={avg_reward:.3f}, solve_rate={solve_rate:.3f}, ' + f'test_pass_rate={avg_pass_rate:.3f}, avg_turns={avg_turns:.1f}') + + # 7. Filter and train (same recipe as the OpenEnv example) + all_input_data: List[Dict[str, Any]] = [] + filtered_old_logps: List[List[float]] = [] + filtered_advantages: List[float] = [] + max_len = rollout_template.max_length or float('inf') + for i, traj in enumerate(all_trajectories): + traj_len = len(traj.get('input_ids') or traj.get('labels') or []) + comp_len = sum(1 for l in (traj.get('labels') or []) if l != -100) + if traj_len > max_len or comp_len == 0: + continue + all_input_data.append(traj) + filtered_old_logps.append(all_old_logps[i]) + filtered_advantages.append(advantages[i]) + + if len(all_input_data) < MODEL_GPUS: + logger.warning(f'[Step {optim_step}] Only {len(all_input_data)} valid trajectories ' + f'after filtering (need >= {MODEL_GPUS}), skipping this batch.') + continue + + total_completions = len(all_input_data) + logger.info(f'[Step {optim_step}] {total_completions}/{n_traj} trajectories ' + f'passed length filter (max_len={max_len})') + + for mb_start in range(0, total_completions, MINI_BATCH_SIZE): + mb_end = min(mb_start + MINI_BATCH_SIZE, total_completions) + model.forward_backward( + inputs=all_input_data[mb_start:mb_end], + old_logps=filtered_old_logps[mb_start:mb_end], + advantages=filtered_advantages[mb_start:mb_end], + micro_batch_size=MICRO_BATCH_SIZE, + ) + model.clip_grad_and_step() + optim_step += 1 + + if optim_step >= MAX_STEPS: + break + if optim_step % SAVE_STEPS == 0: + model.save(f'agentenv-grpo-checkpoint-{optim_step}') + + # 8. Step summary + log_dict = metrics.calculate() + log_dict.update(model.calculate_metric(is_training=True)) + log_dict['avg_turns'] = avg_turns + log_dict['avg_reward'] = avg_reward + log_dict['solve_rate'] = solve_rate + log_dict['test_pass_rate'] = avg_pass_rate + metrics.reset() + logger.info(f'[Step {optim_step}/{MAX_STEPS}] {log_dict}') + + env_pool.shutdown(wait=False) + logger.info(f'Training completed. optim_steps={optim_step}') + model.save('agentenv-grpo-final') + + +if __name__ == '__main__': + main() diff --git a/cookbook/rl/agentenv/agentenv_grpo.sh b/cookbook/rl/agentenv/agentenv_grpo.sh new file mode 100644 index 000000000..478039678 --- /dev/null +++ b/cookbook/rl/agentenv/agentenv_grpo.sh @@ -0,0 +1,54 @@ +#!/bin/sh +set -eu + +# Multi-turn GRPO on MBPP with AgentENV Firecracker sandboxes. +# +# The model writes a Python function, tests it inside a real microVM sandbox and +# submits it; reward is the hidden unit-test pass rate. Sandbox placement / load +# balancing / lifecycle are handled by AgentENV; this script only needs the API +# endpoint. +# +# One-time setup: +# 1. Deploy AgentENV and note its endpoint (server :8000, or gateway :8080). +# 2. Build the sandbox template from the Dockerfile in this folder: +# aenv auth # point at your deployment, any non-empty API key +# aenv build cookbook/rl/agentenv/Dockerfile -t twinkle-code \ +# --cpu-count 1 --memory-mb 1024 +# aenv template watch +# 3. pip install e2b +# +# Run (override anything at invocation): +# AENV_API_URL=http://10.0.0.5:8080 sh agentenv_grpo.sh --max-steps 500 + +# ---- AgentENV connection (client only ever needs this ONE address) ---- +# NOTE: this is THIS SCRIPT's variable name. It is passed as AgentEnv(api_url=...), +# which internally sets the E2B SDK's E2B_API_URL / E2B_SANDBOX_URL (AgentENV +# exposes an E2B-compatible HTTP API). Setting E2B_API_URL alone has NO effect +# here — this script would silently fall back to the default below. +# +# SECURITY: AgentENV has no authorization. Restrict its port to this training +# host with a security group / firewall rule; never expose it to the internet. +export AENV_API_URL="${AENV_API_URL:-http://127.0.0.1:8000}" +export AENV_TEMPLATE="${AENV_TEMPLATE:-twinkle-code}" +# Sandbox idle timeout in seconds; must outlast the slowest rollout in a batch. +export SANDBOX_TIMEOUT="${SANDBOX_TIMEOUT:-600}" +# Concurrent sandbox create/kill HTTP calls from the driver. +export ENV_CONCURRENCY="${ENV_CONCURRENCY:-16}" +# Max tool-calling turns per episode. +export MAX_TURNS="${MAX_TURNS:-6}" + +python agentenv_grpo.py \ + --model-id ms://Qwen/Qwen3.5-4B \ + --model-gpus 4 \ + --sampler-gpus 4 \ + --num-generations 8 \ + --max-tokens 2048 \ + --batch-size 4 \ + --mini-batch-size 8 \ + --micro-batch-size 2 \ + --max-steps 1000 \ + --lr 1e-5 \ + --lora-r 16 \ + --save-steps 500 \ + --adapter-name default \ + "$@" diff --git a/cookbook/rl/agentenv/tools.py b/cookbook/rl/agentenv/tools.py new file mode 100644 index 000000000..7cf3d4bd6 --- /dev/null +++ b/cookbook/rl/agentenv/tools.py @@ -0,0 +1,137 @@ +"""Tool definitions for AgentENV-backed code-writing RL (MBPP). + +Mirrors ``cookbook/rl/openenv_code/tools.py`` — identical task and tool names — +but the tools execute inside a real Firecracker microVM instead of a remote +OpenEnv interpreter session. The differences that matter to the prompt: + + * Full CPython, so ``assert`` / ``try`` / imports / files all work, and the + hidden tests can be replayed as one ordinary script. + * Each ``run_python`` call is a FRESH process, so snippets must be + self-contained (an OpenEnv session, by contrast, keeps its namespace). + +Two tools are exposed to the model: + * ``run_python`` — write a snippet to /workspace/scratch.py and run it. + * ``submit_solution`` — hand in the final function (recorded client-side). +""" +import textwrap +from typing import Any, Dict, List, Tuple + +from twinkle_agentic.envs import AgentEnv + +SYSTEM_PROMPT = """You are an expert Python programmer with access to a Linux sandbox. + +Write a function that solves the given task, verify it, then submit it. + +Rules: +- Use `run_python` to try out your function. Each call runs in a FRESH process, + so every snippet must be self-contained (include the imports and the function + definition, then call it) and must `print(...)` what you want to see. +- The full Python standard library is available, plus `numpy` and `sympy`. +- Match the function name and signature implied by the task description and + the example call, otherwise the hidden tests cannot find your function. +- When your function works, call `submit_solution` with the COMPLETE final + source (imports plus the function definition). +- After submitting, reply with one short sentence and do NOT call any more tools. +- You have a limited number of turns, so do not run redundant code.""" + +TOOL_SCHEMA: List[Dict[str, Any]] = [ + { + 'type': 'function', + 'function': { + 'name': 'run_python', + 'description': 'Run a self-contained Python snippet in the sandbox and return its stdout/stderr.', + 'parameters': { + 'type': 'object', + 'properties': { + 'code': { + 'type': 'string', + 'description': 'Python source to execute. Must print what you want to inspect.', + }, + }, + 'required': ['code'], + }, + }, + }, + { + 'type': 'function', + 'function': { + 'name': 'submit_solution', + 'description': 'Submit the final solution source and end the coding phase.', + 'parameters': { + 'type': 'object', + 'properties': { + 'code': { + 'type': 'string', + 'description': 'Complete final source: imports plus the function definition.', + }, + }, + 'required': ['code'], + }, + }, + }, +] + + +def _run_python(env: AgentEnv, arguments: Dict[str, Any]) -> str: + """Write the snippet to a file and execute it, avoiding shell quoting issues.""" + code = arguments.get('code') + if not code: + return "Error: 'code' argument is required." + env.sandbox.files.write('/workspace/scratch.py', code) + return env.run_command({'command': 'python /workspace/scratch.py', 'cwd': '/workspace'}) + + +def _submit_solution(env: AgentEnv, arguments: Dict[str, Any]) -> str: + """Record the solution on the env; the training loop scores it later.""" + code = (arguments.get('code') or '').strip() + if not code: + return "Error: 'code' argument is required." + env.submitted_code = code + return 'Solution submitted.' + + +def register_tools(env: AgentEnv) -> AgentEnv: + """Attach the task tools to a fresh AgentEnv instance.""" + env.submitted_code = None + return (env.register_tool(TOOL_SCHEMA[0], _run_python).register_tool(TOOL_SCHEMA[1], _submit_solution)) + + +def _build_test_script(solution: str, test_list: List[str], setup_code: str) -> str: + """Build a script that runs each assertion independently and prints a tally. + + Every test is wrapped in its own ``try`` so that one failing assertion (or + one that raises) does not hide the rest — the reward uses the pass rate. + """ + parts = [solution, ''] + if setup_code: + parts += [setup_code, ''] + parts.append('_passed = 0') + for test in test_list: + body = textwrap.indent(test.strip(), ' ') + parts += ['try:', body, ' _passed += 1', 'except Exception:', ' pass'] + parts.append(f"print('TESTS_PASSED', _passed, {len(test_list)})") + return '\n'.join(parts) + '\n' + + +def run_tests(env: AgentEnv, test_list: List[str], setup_code: str = '') -> Tuple[int, int]: + """Replay the hidden tests against the submitted solution inside the sandbox. + + Returns: + ``(n_passed, n_total)``. ``(0, n)`` when nothing was submitted, or when + the script never reaches its tally line (syntax error, timeout, ...). + """ + total = len(test_list) + solution = getattr(env, 'submitted_code', None) + if not solution: + return 0, total + + script = _build_test_script(solution, test_list, setup_code) + env.sandbox.files.write('/workspace/run_tests.py', script) + output = env.run_command({'command': 'python /workspace/run_tests.py', 'cwd': '/workspace'}) + + for line in reversed(output.splitlines()): + if line.startswith('TESTS_PASSED'): + fields = line.split() + if len(fields) >= 3 and fields[1].isdigit(): + return int(fields[1]), total + return 0, total diff --git a/cookbook/rl/code_rl/Dockerfile b/cookbook/rl/code_rl/Dockerfile new file mode 100644 index 000000000..14f20f022 --- /dev/null +++ b/cookbook/rl/code_rl/Dockerfile @@ -0,0 +1,26 @@ +# Sandbox template for cookbook/rl/code_rl with CODE_RL_BACKEND=agentenv +# +# A general-purpose Python code interpreter: the model writes functions, runs +# them, and the training loop replays hidden unit tests inside this sandbox. +# Unlike OpenEnv's AST-interpreter backend, this is a real CPython process in a +# Firecracker microVM, so the full language (assert / try / subprocess / files) +# and arbitrary pip packages are available. +# +# Build this template ONCE against your AgentENV deployment (offline step, +# requires the aenv CLI authenticated via `aenv auth`): +# +# aenv build cookbook/rl/code_rl/Dockerfile -t twinkle-code \ +# --cpu-count 1 --memory-mb 1024 +# aenv template watch # wait until the build is ready +# +# aenv build understands FROM / RUN / ENV / WORKDIR / USER instructions and +# bakes the result into a snapshot, so sandboxes boot in ~50ms with all +# dependencies pre-installed. Rebuild only when this file changes. +FROM python:3.11-slim + +# Use a pip mirror if your sandbox network requires it, e.g.: +# ENV PIP_INDEX_URL=https://mirrors.aliyun.com/pypi/simple +RUN pip install --no-cache-dir numpy sympy + +ENV PYTHONUNBUFFERED=1 +WORKDIR /workspace diff --git a/cookbook/rl/code_rl/README.md b/cookbook/rl/code_rl/README.md new file mode 100644 index 000000000..b5c8c57eb --- /dev/null +++ b/cookbook/rl/code_rl/README.md @@ -0,0 +1,264 @@ +# Code RL (MBPP) + +一份训练脚本，两种执行后端。[English](#english) | 中文 + +## 任务 + +MBPP 数据集，模型写 Python 函数。每条轨迹：模型调 `run_python` 在沙箱里试代码，调 `submit_solution` 提交，训练侧用隐藏测试打分。 + +reward：全部通过 `1.0`；提交了但没全对 `0.1 + 0.4 × 通过率`；没提交 `0.0`。 + +Qwen3.5-4B + LoRA，GRPO，最多 6 轮工具调用。 + +## 两个后端 + +| | `openenv` | `agentenv` | +|---|---|---| +| 代码在哪跑 | OpenEnv 服务的一个 session | 一条轨迹一个 Firecracker microVM | +| 解释器 | AST 解释器（`coding_env`） | 真 CPython | +| 单 env 内存 | KB 级 | ~1GB | +| `run_python` 间状态 | 保留 | 每次新进程 | +| 文件 / 网络 / pip | 不支持 | 支持 | +| async / 装饰器 | 不支持（装饰器静默失效） | 支持 | + +MBPP 用 `openenv` 够。需要 `unittest` + `@patch`、写文件、装包时用 `agentenv`。 + +## 资源需求 + +训练机：8 卡（`--model-gpus 4` + `--sampler-gpus 4`）。显存占用待实测补充。 + +环境机： + +| 后端 | 要求 | +|---|---| +| `openenv` | 普通 CPU 机器，2 核 4G。可与训练机同机 | +| `agentenv` | 裸金属或支持嵌套虚拟化的实例，需 `/dev/kvm`、内核 6.8+、Ubuntu 24.04。容器实例通常不满足 | + +`agentenv` 内存 = `batch-size × num-generations × 1GB + 8GB`，默认 32 并发约 40GB。 + +## 运行：openenv + +环境机： + +```bash +pip install openenv + +# coding_env 是 OpenEnv 仓库里的子包（包名 openenv-coding_env），不在 PyPI 上， +# 只能从源码装。serve.sh 起的 server_app.py 从它 import PythonCodeActEnv 和 +# PyExecutor；它的依赖里带 smolagents，也就是上表那个 AST 解释器的实现。 +git clone https://github.com/huggingface/OpenEnv.git +pip install -e OpenEnv/envs/coding_env + +HOST=127.0.0.1 sh serve.sh +``` + +训练机： + +```bash +pip install openenv +sh run_openenv.sh +``` + +跨机时环境侧绑内网 IP（`HOST=10.0.1.20 sh serve.sh`），训练侧 `OPENENV_BASE_URL=http://10.0.1.20:8000 sh run_openenv.sh`。 + +## 运行：agentenv + +环境机装服务端（Ubuntu 24.04，同时装上 `aenv` CLI）： + +```bash +curl -fsSL https://raw.githubusercontent.com/kvcache-ai/AgentENV/main/scripts/install.sh | sudo bash +sudo systemctl start aenv +``` + +单独装 CLI（走 Docker 部署，或 CLI 与服务端不同机时）： + +```bash +curl -fsSL https://raw.githubusercontent.com/kvcache-ai/AgentENV/main/scripts/install-cli.sh | bash +``` + +认证并构建模板（一次，改 `Dockerfile` 才需重建）： + +```bash +aenv auth # URL 填服务端地址；AgentENV 无认证，API key 随便填 +aenv build cookbook/rl/code_rl/Dockerfile -t twinkle-code --cpu-count 1 --memory-mb 1024 +aenv template watch +``` + +训练机： + +```bash +pip install e2b +AENV_API_URL=http://<环境机IP>:8000 sh run_agentenv.sh +``` + +## 参数 + +命令行参数透传给 `train.py`，覆盖 `common_args.sh` 里的默认值： + +```bash +sh run_openenv.sh --max-steps 500 --batch-size 8 +``` + +冒烟测试（`batch-size × num-generations` 需 ≥ `--model-gpus`）： + +```bash +sh run_openenv.sh --batch-size 2 --num-generations 4 --max-steps 2 +``` + +## 实验记录 + +待补充。 + +| backend | 配置 | steps | reward | 备注 | +|---|---|---|---|---| +| | | | | | + +## 文件 + +| 文件 | 作用 | +|---|---| +| `train.py` | 训练逻辑，与后端无关 | +| `backends/openenv.py` `backends/agentenv.py` | env 构造、提示词、工具、隐藏测试回放 | +| `common_args.sh` | 训练超参 | +| `run_openenv.sh` `run_agentenv.sh` | 启动命令 | +| `serve.sh` `server_app.py` | OpenEnv 服务端 | +| `Dockerfile` | AgentENV 沙箱模板 | + +加后端：写 `backends/xxx.py`（`NAME`、`SYSTEM_PROMPT`、`TOOL_SCHEMA`、`make_env()`、`run_tests()`、`describe()`）和 `run_xxx.sh`，不改 `train.py`。 + +跨网络部署见 `docs/source_zh/使用指引/Agentic RL部署与训练.md`。 + +--- + + + +# Code RL (MBPP) — English + +One training script, two execution backends. + +## Task + +MBPP dataset, the model writes Python functions. Per trajectory: the model calls `run_python` to try code in a sandbox, calls `submit_solution`, and the trainer scores it against hidden tests. + +Reward: `1.0` if all tests pass; `0.1 + 0.4 × pass_rate` if submitted but incorrect; `0.0` if never submitted. + +Qwen3.5-4B + LoRA, GRPO, up to 6 tool-calling turns. + +## Two backends + +| | `openenv` | `agentenv` | +|---|---|---| +| Where code runs | A session on an OpenEnv server | One Firecracker microVM per trajectory | +| Interpreter | AST interpreter (`coding_env`) | Real CPython | +| Memory per env | KBs | ~1GB | +| State between `run_python` calls | Persists | Fresh process each call | +| Files / network / pip | No | Yes | +| async / decorators | No (decorators silently ignored) | Yes | + +`openenv` is enough for MBPP. Use `agentenv` when you need `unittest` + `@patch`, file writes, or pip installs. + +## Resources + +Training host: 8 GPUs (`--model-gpus 4` + `--sampler-gpus 4`). Measured memory usage to be filled in. + +Environment host: + +| Backend | Requirement | +|---|---| +| `openenv` | An ordinary CPU machine, 2 cores / 4GB. Can be the training host itself | +| `agentenv` | Bare metal or nested-virtualisation instance; needs `/dev/kvm`, kernel 6.8+, Ubuntu 24.04. Container instances usually do not qualify | + +`agentenv` memory = `batch-size × num-generations × 1GB + 8GB`, i.e. ~40GB at the default 32 concurrent sandboxes. + +## Run: openenv + +Environment host: + +```bash +pip install openenv + +# coding_env is a sub-package inside the OpenEnv repo (distribution name +# openenv-coding_env). It is not on PyPI, so it has to be installed from source. +# server_app.py, which serve.sh runs, imports PythonCodeActEnv and PyExecutor +# from it; its dependencies pull in smolagents — the AST interpreter above. +git clone https://github.com/huggingface/OpenEnv.git +pip install -e OpenEnv/envs/coding_env + +HOST=127.0.0.1 sh serve.sh +``` + +Training host: + +```bash +pip install openenv +sh run_openenv.sh +``` + +Across hosts, bind the environment to its private NIC (`HOST=10.0.1.20 sh serve.sh`) and run `OPENENV_BASE_URL=http://10.0.1.20:8000 sh run_openenv.sh`. + +## Run: agentenv + +Install the server on the environment host (Ubuntu 24.04; this also installs the `aenv` CLI): + +```bash +curl -fsSL https://raw.githubusercontent.com/kvcache-ai/AgentENV/main/scripts/install.sh | sudo bash +sudo systemctl start aenv +``` + +Install the CLI alone (Docker deployment, or CLI on a different machine): + +```bash +curl -fsSL https://raw.githubusercontent.com/kvcache-ai/AgentENV/main/scripts/install-cli.sh | bash +``` + +Authenticate and build the template (once; rebuild only when `Dockerfile` changes): + +```bash +aenv auth # URL is the server address; AgentENV has no authorization, any API key works +aenv build cookbook/rl/code_rl/Dockerfile -t twinkle-code --cpu-count 1 --memory-mb 1024 +aenv template watch +``` + +Training host: + +```bash +pip install e2b +AENV_API_URL=http://:8000 sh run_agentenv.sh +``` + +## Arguments + +Command-line arguments are forwarded to `train.py` and override the defaults in `common_args.sh`: + +```bash +sh run_openenv.sh --max-steps 500 --batch-size 8 +``` + +Smoke test (`batch-size × num-generations` must be ≥ `--model-gpus`): + +```bash +sh run_openenv.sh --batch-size 2 --num-generations 4 --max-steps 2 +``` + +## Results + +To be filled in. + +| backend | config | steps | reward | notes | +|---|---|---|---|---| +| | | | | | + +## Files + +| File | Role | +|---|---| +| `train.py` | Training logic, backend-agnostic | +| `backends/openenv.py`, `backends/agentenv.py` | env construction, prompt, tools, hidden-test replay | +| `common_args.sh` | Training hyper-parameters | +| `run_openenv.sh`, `run_agentenv.sh` | Launch commands | +| `serve.sh`, `server_app.py` | The OpenEnv server | +| `Dockerfile` | The AgentENV sandbox template | + +To add a backend: write `backends/xxx.py` (`NAME`, `SYSTEM_PROMPT`, `TOOL_SCHEMA`, `make_env()`, `run_tests()`, `describe()`) and a `run_xxx.sh`. `train.py` stays untouched. + +For cross-network deployment see `docs/source_en/Usage Guide/Agentic-RL-Deployment.md`. diff --git a/cookbook/rl/code_rl/backends/__init__.py b/cookbook/rl/code_rl/backends/__init__.py new file mode 100644 index 000000000..cb0ede971 --- /dev/null +++ b/cookbook/rl/code_rl/backends/__init__.py @@ -0,0 +1,51 @@ +"""Execution backends for the code-writing RL task. + +``train.py`` is backend-agnostic: it owns the dataset, the GRPO loop and the +reward shaping, and delegates everything that depends on *where code runs* to +one of the modules here. + +A backend module must expose: + +``NAME`` + Short identifier used in logs and checkpoint names. +``SYSTEM_PROMPT`` + Task prompt. This is NOT shared, because the execution semantics differ in + ways the model must know about — most importantly whether state survives + between ``run_python`` calls, and which modules are importable. A prompt + that describes the wrong backend sends the policy down the wrong path. +``TOOL_SCHEMA`` + OpenAI-format tool list advertised to the model. Both backends use the same + two tool *names* (``run_python``, ``submit_solution``) so trajectories stay + comparable; only the descriptions differ. +``make_env() -> Env`` + Build one environment for one trajectory, with tools registered. Must set + ``env.submitted_code = None``. +``run_tests(env, test_list, setup_code) -> tuple[int, int]`` + Replay the hidden tests against ``env.submitted_code`` and return + ``(n_passed, n_total)``. Returns ``(0, n)`` when nothing was submitted. +``describe() -> str`` + One line of connection info for the startup log. + +Env lifecycle (``reset``/``close``) is part of the ``Env`` interface, so +``train.py`` calls those directly. +""" +from importlib import import_module +from types import ModuleType + +AVAILABLE = ('openenv', 'agentenv') + + +def get_backend(name: str) -> ModuleType: + """Import a backend module by short name. + + Args: + name: One of ``AVAILABLE``. + + Raises: + ValueError: On an unknown name, listing the valid ones — a typo in + ``CODE_RL_BACKEND`` should fail immediately rather than silently + falling back to a default and training the wrong thing. + """ + if name not in AVAILABLE: + raise ValueError(f'Unknown backend {name!r}. Available: {", ".join(AVAILABLE)}') + return import_module(f'backends.{name}') diff --git a/cookbook/rl/code_rl/backends/agentenv.py b/cookbook/rl/code_rl/backends/agentenv.py new file mode 100644 index 000000000..0c16388d5 --- /dev/null +++ b/cookbook/rl/code_rl/backends/agentenv.py @@ -0,0 +1,156 @@ +"""AgentENV backend: one Firecracker microVM per trajectory. + +Code runs in real CPython inside a microVM. Two properties shape the prompt +below, and both are the OPPOSITE of the OpenEnv backend: + + * Each ``run_python`` call is a FRESH process, so every snippet must be + self-contained — nothing carries over between calls. + * It is a full interpreter: the whole standard library works, decorators and + ``unittest`` behave normally, and files can be written. +""" +import os +import textwrap +from typing import Any, Dict, List, Tuple + +from twinkle_agentic.envs import AgentEnv + +NAME = 'agentenv' + +API_URL = os.environ.get('AENV_API_URL', 'http://127.0.0.1:8000') +TEMPLATE = os.environ.get('AENV_TEMPLATE', 'twinkle-code') +SANDBOX_TIMEOUT = int(os.environ.get('SANDBOX_TIMEOUT', '600')) +COMMAND_TIMEOUT = int(os.environ.get('AENV_COMMAND_TIMEOUT', '60')) + +SYSTEM_PROMPT = """You are an expert Python programmer with access to a Linux sandbox. + +Solve the task by writing a Python function. + +- Use `run_python` to try out your function. Each call runs in a FRESH process, + so every snippet must be self-contained (include the imports and the function + definition, then call it) and must `print(...)` what you want to see. +- The full Python standard library is available, plus `numpy` and `sympy`. +- When you are confident, call `submit_solution` with the complete final source + (imports plus the function definition). + +Submit exactly once, and only after the code runs correctly.""" + +TOOL_SCHEMA: List[Dict[str, Any]] = [ + { + 'type': 'function', + 'function': { + 'name': 'run_python', + 'description': 'Run a self-contained Python snippet in the sandbox and return its stdout/stderr.', + 'parameters': { + 'type': 'object', + 'properties': { + 'code': { + 'type': 'string', + 'description': 'Python source to execute. Must print what you want to inspect.', + }, + }, + 'required': ['code'], + }, + }, + }, + { + 'type': 'function', + 'function': { + 'name': 'submit_solution', + 'description': 'Submit the final solution source and end the coding phase.', + 'parameters': { + 'type': 'object', + 'properties': { + 'code': { + 'type': 'string', + 'description': 'Complete final source: imports plus the function definition.', + }, + }, + 'required': ['code'], + }, + }, + }, +] + + +def _run_python(env: AgentEnv, arguments: Dict[str, Any]) -> str: + """Write the snippet to a file and execute it, avoiding shell quoting issues.""" + code = arguments.get('code') + if not code: + return "Error: 'code' argument is required." + env.sandbox.files.write('/workspace/scratch.py', code) + return env.run_command({'command': 'python /workspace/scratch.py', 'cwd': '/workspace'}) + + +def _submit_solution(env: AgentEnv, arguments: Dict[str, Any]) -> str: + """Record the solution on the env; the training loop scores it later.""" + code = (arguments.get('code') or '').strip() + if not code: + return "Error: 'code' argument is required." + env.submitted_code = code + return 'Solution submitted.' + + +def make_env() -> AgentEnv: + """Boot one sandbox per trajectory, exposing only the two task tools. + + ``include_default_tools=False`` hides AgentEnv's built-ins (raw command + execution, file read/write) so the action space matches the task exactly and + reward attribution stays clean. + """ + env = AgentEnv( + template=TEMPLATE, + api_url=API_URL, + sandbox_timeout=SANDBOX_TIMEOUT, + command_timeout=COMMAND_TIMEOUT, + include_default_tools=False, + ) + env.submitted_code = None + return (env.register_tool(TOOL_SCHEMA[0], _run_python).register_tool(TOOL_SCHEMA[1], _submit_solution)) + + +def _build_test_script(solution: str, test_list: List[str], setup_code: str) -> str: + """Build a script that runs each assertion independently and prints a tally. + + Every test is wrapped in its own ``try`` so that one failing assertion (or + one that raises) does not hide the rest — the reward uses the pass rate. + """ + parts = [solution, ''] + if setup_code: + parts += [setup_code, ''] + parts.append('_passed = 0') + for test in test_list: + body = textwrap.indent(test.strip(), ' ') + parts += ['try:', body, ' _passed += 1', 'except Exception:', ' pass'] + parts.append(f"print('TESTS_PASSED', _passed, {len(test_list)})") + return '\n'.join(parts) + '\n' + + +def run_tests(env: AgentEnv, test_list: List[str], setup_code: str = '') -> Tuple[int, int]: + """Replay the hidden tests against the submitted solution inside the sandbox. + + Real CPython means the tests can run as one ordinary script, unlike the + OpenEnv backend which drives them one expression at a time. + + Returns: + ``(n_passed, n_total)``. ``(0, n)`` when nothing was submitted, or when + the script never reaches its tally line (syntax error, timeout, ...). + """ + total = len(test_list) + solution = getattr(env, 'submitted_code', None) + if not solution: + return 0, total + + script = _build_test_script(solution, test_list, setup_code) + env.sandbox.files.write('/workspace/run_tests.py', script) + output = env.run_command({'command': 'python /workspace/run_tests.py', 'cwd': '/workspace'}) + + for line in reversed(output.splitlines()): + if line.startswith('TESTS_PASSED'): + fields = line.split() + if len(fields) >= 3 and fields[1].isdigit(): + return int(fields[1]), total + return 0, total + + +def describe() -> str: + return f'AgentENV microVM: api_url={API_URL}, template={TEMPLATE}' diff --git a/cookbook/rl/code_rl/backends/openenv.py b/cookbook/rl/code_rl/backends/openenv.py new file mode 100644 index 000000000..7e9dbb27d --- /dev/null +++ b/cookbook/rl/code_rl/backends/openenv.py @@ -0,0 +1,153 @@ +"""OpenEnv server-mode backend: one WebSocket session per trajectory. + +Code runs in a remote OpenEnv ``coding_env`` interpreter session. Two properties +shape the prompt below: + + * The session namespace PERSISTS across ``run_python`` calls, so the model can + define a function in one turn and probe it in the next. + * ``coding_env`` executes through an AST interpreter, not real CPython. There + is no file or network access, decorators are silently ignored, and only + authorised modules can be imported (see ``server_app.py``). +""" +import os +from typing import Any, Dict, List, Tuple + +from twinkle_agentic.envs import OpenEnvClient + +NAME = 'openenv' + +BASE_URL = os.environ.get('OPENENV_BASE_URL', 'http://127.0.0.1:8000') +ENV_NAME = os.environ.get('OPENENV_ENV_NAME', 'coding_env') +# Code execution can be slow; keep the per-message timeout generous. Note the +# executor's own caps still apply (operation count, while-loop iterations, and +# a wall-clock limit in newer smolagents releases). +MESSAGE_TIMEOUT_S = float(os.environ.get('OPENENV_MESSAGE_TIMEOUT_S', '120')) + +SYSTEM_PROMPT = """You are an expert Python programmer with access to a Python interpreter. + +Solve the task by writing a Python function. + +- Use `run_python` to define and test your function. The interpreter keeps its + state between calls, so a function defined in one call stays available. +- Available modules: math, re, collections, itertools, functools, operator, + heapq, bisect, string, statistics, fractions, decimal, datetime, copy, json. + There is no file or network access. +- When you are confident, call `submit_solution` with the complete final source + (imports plus the function definition). + +Submit exactly once, and only after the code runs correctly.""" + +TOOL_SCHEMA: List[Dict[str, Any]] = [ + { + 'type': 'function', + 'function': { + 'name': 'run_python', + 'description': 'Execute a Python snippet in the interpreter and return its stdout/stderr.', + 'parameters': { + 'type': 'object', + 'properties': { + 'code': { + 'type': 'string', + 'description': 'Python source to execute. Use print(...) to inspect values.', + }, + }, + 'required': ['code'], + }, + }, + }, + { + 'type': 'function', + 'function': { + 'name': 'submit_solution', + 'description': 'Submit the final solution source and end the coding phase.', + 'parameters': { + 'type': 'object', + 'properties': { + 'code': { + 'type': 'string', + 'description': 'Complete final source: imports plus the function definition.', + }, + }, + 'required': ['code'], + }, + }, + }, +] + + +def _submit_solution(env: OpenEnvClient, arguments: Dict[str, Any]) -> str: + """Record the solution on the env; the training loop scores it later.""" + code = (arguments.get('code') or '').strip() + if not code: + return "Error: 'code' argument is required." + env.submitted_code = code + return 'Solution submitted.' + + +def make_env() -> OpenEnvClient: + """Open one session per trajectory, exposing only the two task tools. + + ``run_python`` stays server-backed: the default action mapper turns + ``{'code': ...}`` straight into the env's ``CodeAction``. Only + ``submit_solution`` needs a client-side handler. + """ + env = OpenEnvClient( + env_name=ENV_NAME, + base_url=BASE_URL, + tools=[TOOL_SCHEMA[0]], + message_timeout_s=MESSAGE_TIMEOUT_S, + ) + env.submitted_code = None + return env.register_tool(TOOL_SCHEMA[1], _submit_solution) + + +def _ok(result) -> bool: + return not (getattr(result.observation, 'exit_code', 0) or 0) + + +def _last_line(text: str) -> str: + lines = [line for line in (text or '').splitlines() if line.strip()] + return lines[-1].strip() if lines else '' + + +def run_tests(env: OpenEnvClient, test_list: List[str], setup_code: str = '') -> Tuple[int, int]: + """Run the hidden tests against the submitted solution in the same session. + + Each assertion is executed as its own ``print()`` step rather than as + one block of ``assert`` statements. Two reasons: printing the boolean tells + "the assertion evaluated to False" apart from "the code blew up", whereas a + failed ``assert`` surfaces as an exception indistinguishable from a crash + inside the solution; and one step per test isolates a test that raises, so + the remaining ones still run. + + The executor does support ``assert``/``try``, so this is a diagnosability + choice, not a capability workaround. + + Returns: + ``(n_passed, n_total)``. ``(0, n)`` when nothing was submitted or the + solution itself fails to execute. + """ + total = len(test_list) + solution = getattr(env, 'submitted_code', None) + if not solution: + return 0, total + + # Define the submitted solution in the session namespace. + if not _ok(env.execute({'code': solution})): + return 0, total + if setup_code and not _ok(env.execute({'code': setup_code})): + return 0, total + + passed = 0 + for test in test_list: + expr = test.strip() + if expr.startswith('assert '): + expr = expr[len('assert '):] + result = env.execute({'code': f'print({expr})'}) + if _ok(result) and _last_line(getattr(result.observation, 'stdout', '')) == 'True': + passed += 1 + return passed, total + + +def describe() -> str: + return f'OpenEnv server mode: base_url={BASE_URL}, env={ENV_NAME}' diff --git a/cookbook/rl/code_rl/common_args.sh b/cookbook/rl/code_rl/common_args.sh new file mode 100644 index 000000000..6e2265b4e --- /dev/null +++ b/cookbook/rl/code_rl/common_args.sh @@ -0,0 +1,36 @@ +# Shared training hyper-parameters for both backends. Sourced by run_*.sh. +# +# Kept in one place so the two backends stay a controlled comparison: if these +# differed, a change in reward could not be attributed to the execution backend. +# Backend-specific settings live in the run_*.sh scripts and in backends/. +# +# Override anything at invocation, e.g.: +# sh run_openenv.sh --max-steps 500 --batch-size 8 + +# Capacity note: the backend must host BATCH_SIZE x NUM_GENERATIONS concurrent +# envs (32 with the defaults below). +# - OpenEnv: WORKERS x MAX_CONCURRENT_ENVS in serve.sh (256 by default) +# - AgentENV: 32 x sandbox memory + ~8GB for the host +# +# Floor: trajectories (batch-size x num-generations) must stay >= --model-gpus, +# or every batch is dropped by the length filter with only a warning. +TRAIN_ARGS=" + --model-id ms://Qwen/Qwen3.5-4B + --model-gpus 4 + --sampler-gpus 4 + --num-generations 8 + --max-tokens 2048 + --batch-size 4 + --mini-batch-size 8 + --micro-batch-size 2 + --max-steps 1000 + --lr 1e-5 + --lora-r 16 + --save-steps 500 + --adapter-name default +" + +# Max tool-calling turns per episode. +export MAX_TURNS="${MAX_TURNS:-6}" +# Concurrent reset/score calls issued from the driver. +export ENV_CONCURRENCY="${ENV_CONCURRENCY:-16}" diff --git a/cookbook/rl/code_rl/run_agentenv.sh b/cookbook/rl/code_rl/run_agentenv.sh new file mode 100644 index 000000000..e8cf6c9ef --- /dev/null +++ b/cookbook/rl/code_rl/run_agentenv.sh @@ -0,0 +1,45 @@ +#!/bin/sh +# Code-writing GRPO on AgentENV (one Firecracker microVM per trajectory). +# +# Prerequisites (once, offline): +# 1. Deploy AgentENV (single server, or gateway+scheduler cluster). +# Needs /dev/kvm, kernel 6.8+, CAP_SYS_ADMIN — container instances usually +# do not qualify; see the sandbox doc in docs. +# 2. Build the sandbox template from the Dockerfile in this folder: +# aenv build cookbook/rl/code_rl/Dockerfile -t twinkle-code \ +# --cpu-count 1 --memory-mb 1024 +# 3. On the training host: pip install e2b +# +# Memory budget: BATCH_SIZE x NUM_GENERATIONS sandboxes run at once, so the +# defaults (32) need roughly 32 x 1GB + 8GB for AgentENV and the OS. +# +# Run (override anything at invocation): +# sh run_agentenv.sh +# AENV_API_URL=http://10.0.0.5:8000 sh run_agentenv.sh --max-steps 500 +set -eu +cd "$(dirname "$0")" + +export CODE_RL_BACKEND=agentenv + +# ---- AgentENV connection (the client only ever needs this ONE address) ---- +# Passed through as AgentEnv(api_url=...), which internally sets the E2B SDK's +# E2B_API_URL / E2B_SANDBOX_URL (AgentENV exposes an E2B-compatible HTTP API). +# Setting E2B_API_URL alone has NO effect here — this script would silently fall +# back to the default below. +# +# SECURITY: AgentENV has no authorization. Restrict its port to this training +# host with a security group / firewall rule; never expose it to the internet. +export AENV_API_URL="${AENV_API_URL:-http://127.0.0.1:8000}" +# Template built in step 2 above. +export AENV_TEMPLATE="${AENV_TEMPLATE:-twinkle-code}" +# Sandbox lifetime; must outlast one episode plus the test replay. +export SANDBOX_TIMEOUT="${SANDBOX_TIMEOUT:-600}" +# Per-command timeout inside the sandbox. +export AENV_COMMAND_TIMEOUT="${AENV_COMMAND_TIMEOUT:-60}" + +. ./common_args.sh + +echo "Backend: agentenv | api: $AENV_API_URL | template: $AENV_TEMPLATE" + +# shellcheck disable=SC2086 +python train.py $TRAIN_ARGS "$@" diff --git a/cookbook/rl/code_rl/run_openenv.sh b/cookbook/rl/code_rl/run_openenv.sh new file mode 100644 index 000000000..04a77e61e --- /dev/null +++ b/cookbook/rl/code_rl/run_openenv.sh @@ -0,0 +1,34 @@ +#!/bin/sh +# Code-writing GRPO on a remote OpenEnv server (one session per trajectory). +# +# Prerequisites: +# 1. On the environment host: +# pip install openenv && pip install -e /path/to/OpenEnv/envs/coding_env +# HOST=127.0.0.1 sh serve.sh # see serve.sh for network options +# 2. On the training host: +# pip install openenv +# +# Run (override anything at invocation): +# sh run_openenv.sh +# OPENENV_BASE_URL=http://10.0.0.5:8000 sh run_openenv.sh --max-steps 500 +set -eu +cd "$(dirname "$0")" + +export CODE_RL_BACKEND=openenv + +# ---- OpenEnv connection (a load-balancer address works here too) ---- +# For a cross-network setup this is the LOCAL end of an SSH port forward +# (http://127.0.0.1:8000) rather than the remote address — see the deployment +# guide in docs. +export OPENENV_BASE_URL="${OPENENV_BASE_URL:-http://127.0.0.1:8000}" +# Environment package providing the client + Action classes. +export OPENENV_ENV_NAME="${OPENENV_ENV_NAME:-coding_env}" +# Per-message timeout. The executor has its own caps on top of this. +export OPENENV_MESSAGE_TIMEOUT_S="${OPENENV_MESSAGE_TIMEOUT_S:-120}" + +. ./common_args.sh + +echo "Backend: openenv | server: $OPENENV_BASE_URL | env: $OPENENV_ENV_NAME" + +# shellcheck disable=SC2086 +python train.py $TRAIN_ARGS "$@" diff --git a/cookbook/rl/code_rl/serve.sh b/cookbook/rl/code_rl/serve.sh new file mode 100644 index 000000000..eeb8f349c --- /dev/null +++ b/cookbook/rl/code_rl/serve.sh @@ -0,0 +1,52 @@ +#!/bin/sh +set -eu + +# Start the OpenEnv code environment server for cookbook/rl/code_rl. +# +# This is the REMOTE side: run it on any machine reachable from the training +# driver — no Docker, no KVM, no GPU. The training script only needs the +# resulting URL. +# +# One-time setup: +# pip install openenv +# pip install -e /path/to/OpenEnv/envs/coding_env +# +# Usage: +# sh serve.sh # 4 workers x 64 sessions = 256 sessions +# WORKERS=8 PORT=9000 sh serve.sh +# HOST=127.0.0.1 sh serve.sh # same-machine training: no network exposure +# +# SECURITY: OpenEnv has no authentication. Anyone who can reach this port can +# execute code in the sandbox and consume your capacity. Either bind to +# 127.0.0.1 (when training runs on this same host), or restrict the port to the +# training host's IP with a security group / firewall rule. Never leave the +# default 0.0.0.0 reachable from the public internet. +# +# Then point the training script at it: +# OPENENV_BASE_URL=http://:8000 sh run_openenv.sh + +HOST="${HOST:-0.0.0.0}" +PORT="${PORT:-8000}" +# Uvicorn worker processes. Each worker holds its own app instance, so total +# capacity is WORKERS x MAX_CONCURRENT_ENVS. Size it above the number of +# concurrent trajectories (BATCH_SIZE x NUM_GENERATIONS), otherwise the extra +# WebSocket connections are rejected at capacity. +WORKERS="${WORKERS:-4}" +export MAX_CONCURRENT_ENVS="${MAX_CONCURRENT_ENVS:-64}" + +cd "$(dirname "$0")" + +echo "Serving twinkle_code_env on ${HOST}:${PORT}" +echo " workers=${WORKERS}, max_concurrent_envs=${MAX_CONCURRENT_ENVS} per worker" +echo " capacity=$((WORKERS * MAX_CONCURRENT_ENVS)) concurrent sessions" +if [ "$HOST" = "0.0.0.0" ]; then + echo " WARNING: bound to 0.0.0.0 with no authentication — restrict this port" + echo " to the training host, or use HOST=127.0.0.1 if training is local." +fi + +exec uvicorn server_app:app --host "$HOST" --port "$PORT" --workers "$WORKERS" + +# Docker alternative (build coding_env's image, then override the app module): +# docker build -t coding-env:latest -f /path/to/OpenEnv/envs/coding_env/server/Dockerfile /path/to/OpenEnv +# Note the image serves upstream coding_env.server.app, which is capped at one +# session; mount this file and point uvicorn at it to keep the concurrency fix. diff --git a/cookbook/rl/code_rl/server_app.py b/cookbook/rl/code_rl/server_app.py new file mode 100644 index 000000000..2b2a000d7 --- /dev/null +++ b/cookbook/rl/code_rl/server_app.py @@ -0,0 +1,113 @@ +"""OpenEnv server for the code-writing RL task — server mode, no Docker needed. + +Wraps OpenEnv's ``PythonCodeActEnv`` (``OpenEnv/envs/coding_env``) with the +three changes a training workload needs. Each one is a deliberate deviation +from upstream defaults: + +1. ``SUPPORTS_CONCURRENT_SESSIONS = True``. Upstream leaves this at the + conservative default, which caps the server at ONE session + (``create_app`` raises ``ConcurrencyConfigurationError`` for + ``max_concurrent_envs > 1`` otherwise). The class is in fact session-safe: + ``__init__`` builds a private executor and state and shares nothing, and + ``create_app`` receives the class as a *factory*, so every WebSocket + connection gets a fresh instance. + +2. A wider import whitelist. smolagents' ``LocalPythonExecutor`` authorises + only ``json`` by default, so ``import math`` / ``collections`` — which many + MBPP solutions need — would fail. + +3. No reward transforms. ``coding_env.create_safe_coding_transform()`` + overwrites ``observation.reward`` with code-style heuristics (-1.0 when the + code matches ``open(`` / ``import os``, +0.1 for short code). For this task + the reward must come from unit tests, computed by the training script, so a + style score on the same channel is noise. + +Prerequisites:: + + pip install openenv + pip install -e /path/to/OpenEnv/envs/coding_env # brings in smolagents + +Run:: + + sh serve.sh + # or explicitly: + MAX_CONCURRENT_ENVS=64 uvicorn server_app:app --host 0.0.0.0 --port 8000 --workers 4 + +Environment variables: + MAX_CONCURRENT_ENVS: concurrent sessions per worker process (default 64). +""" +import os + +from coding_env.models import CodeAction, CodeObservation +from coding_env.server.python_codeact_env import PythonCodeActEnv +from coding_env.server.python_executor import PyExecutor +from openenv.core.env_server import create_app + +# Modules the sandboxed executor may import. Keep this list tight: it is the +# only thing standing between model-generated code and this process, since +# LocalPythonExecutor is an AST interpreter, not an OS-level sandbox. +ALLOWED_IMPORTS = [ + 'math', + 're', + 'collections', + 'itertools', + 'functools', + 'operator', + 'heapq', + 'bisect', + 'string', + 'statistics', + 'fractions', + 'decimal', + 'datetime', + 'copy', + 'json', +] + +MAX_CONCURRENT_ENVS = int(os.environ.get('MAX_CONCURRENT_ENVS', '64')) + + +class ConcurrentCodeEnv(PythonCodeActEnv): + """Session-isolated Python executor with a task-appropriate config.""" + + SUPPORTS_CONCURRENT_SESSIONS = True + + def __init__(self): + super().__init__() + self._configure() + + def reset(self, **kwargs): + # The parent's reset() rebuilds the executor and the transform with + # upstream defaults, so re-apply our config afterwards. **kwargs + # absorbs the seed / episode_id the server forwards from the client. + observation = super().reset() + self._configure() + return observation + + def _configure(self) -> None: + self._executor = PyExecutor(additional_imports=list(ALLOWED_IMPORTS)) + # Drop the style/safety reward heuristics; reward comes from the tests. + self.transform = None + + +app = create_app( + ConcurrentCodeEnv, + CodeAction, + CodeObservation, + env_name='twinkle_code_env', + max_concurrent_envs=MAX_CONCURRENT_ENVS, +) + + +def main(): + """Entry point for ``python server_app.py`` (single worker).""" + import uvicorn + uvicorn.run( + app, + host=os.environ.get('HOST', '0.0.0.0'), + port=int(os.environ.get('PORT', '8000')), + ) + + +if __name__ == '__main__': + main() diff --git a/cookbook/rl/code_rl/train.py b/cookbook/rl/code_rl/train.py new file mode 100644 index 000000000..28761f024 --- /dev/null +++ b/cookbook/rl/code_rl/train.py @@ -0,0 +1,408 @@ +"""Code-writing RL (MBPP) with GRPO — one training script, two execution backends. + +The task, the reward and the GRPO loop live here and are identical regardless of +where code runs. Everything backend-specific (env construction, tool handlers, +system prompt, hidden-test replay) lives in ``backends/``; see that package's +docstring for the contract. + +Pick a backend with ``CODE_RL_BACKEND``: + + openenv One WebSocket session per trajectory on a remote OpenEnv server. + Lightweight (KBs per env), AST interpreter, no files/network. + agentenv One Firecracker microVM per trajectory. Real CPython (~1GB per env), + so decorators, unittest+mock, files and pip all work. + +How the loop works: + 1. make_env() per trajectory, reset() in parallel (each reset is a blocking + network call — a batch of 32 must not serialize). + 2. MultiTurnRollout drives tool calls: run_python explores, submit_solution + records the final source. + 3. After the rollout, the hidden MBPP tests are replayed against the submitted + code to compute the reward. + 4. GRPO advantages are group-relative across NUM_GENERATIONS rollouts of the + same problem. + +No EnvPool / @remote_class: in both backends the session or sandbox lives +server-side, so sharding envs through Ray buys nothing and only adds an RPC hop. + +Run: + sh run_openenv.sh # or: sh run_agentenv.sh + OPENENV_BASE_URL=http://10.0.0.5:8000 sh run_openenv.sh --max-steps 500 +""" +import os +import re +from concurrent.futures import ThreadPoolExecutor +from typing import Any, Dict, List, Tuple + +import swanlab +import torch +from peft import LoraConfig + +import twinkle +from backends import get_backend +from twinkle import DeviceGroup, DeviceMesh, get_device_placement, get_logger +from twinkle.advantage import GRPOAdvantage +from twinkle.checkpoint_engine import CheckpointEngineManager +from twinkle.cli import CLI +from twinkle.data_format import SamplingParams +from twinkle.metric import CompletionRewardMetric +from twinkle.model import TransformersModel +from twinkle.processor import InputProcessor +from twinkle.sampler import vLLMSampler +from twinkle.template import Qwen3_5Template +from twinkle_agentic.envs import EnvTool +from twinkle_agentic.rollout.multi_turn import MultiTurnRollout +from twinkle_agentic.tools.tool_manager import ToolManager + +logger = get_logger() +args = CLI.from_args() + +# ========== Backend selection ========== +backend = get_backend(os.environ.get('CODE_RL_BACKEND', 'openenv')) + +# ========== Configuration ========== +MODEL_ID = args.model.model_id or 'ms://Qwen/Qwen3.5-4B' + +MODEL_GPUS = args.infra.model_gpus or 4 +SAMPLER_GPUS = args.infra.sampler_gpus or 4 +NUM_GPUS = MODEL_GPUS + SAMPLER_GPUS + +NUM_GENERATIONS = args.rl.num_generations or 8 +MAX_NEW_TOKENS = args.sampling.max_tokens or 2048 +LEARNING_RATE = args.optimizer.learning_rate or 1e-5 +MAX_STEPS = args.training.max_steps or 1000 +BATCH_SIZE = args.training.batch_size or 4 +MINI_BATCH_SIZE = args.training.mini_batch_size or 8 +MICRO_BATCH_SIZE = args.training.micro_batch_size or 2 +GRADIENT_ACCUMULATION_STEPS = args.training.gradient_accumulation_steps or 1 +ADAPTER_NAME = args.lora.adapter_name or 'default' +SAVE_STEPS = args.training.save_steps or 500 +LORA_RANK = args.lora.lora_r or 16 +MAX_TURNS = int(os.environ.get('MAX_TURNS', '6')) + +# Parallelism for the blocking reset/score calls in the driver. +ENV_CONCURRENCY = int(os.environ.get('ENV_CONCURRENCY', '16')) + + +# ========== Dataset (MBPP) ========== +def load_mbpp() -> List[Dict[str, Any]]: + """Load MBPP as [{'prompt', 'test_list', 'test_setup_code'}].""" + from modelscope.msdatasets import MsDataset + raw = MsDataset.load('opencompass/mbpp', subset_name='full', split='train') + samples = [] + for row in raw: + test_list = list(row['test_list'] or []) + if not test_list: + continue + # The first assertion is handed to the model as a signature hint: MBPP + # descriptions alone do not pin the function name, and the hidden tests + # call it by name. + prompt = (f"{row['text']}\n\n" + f'Your function must satisfy this call:\n{test_list[0]}') + samples.append({ + 'prompt': prompt, + 'test_list': test_list, + 'test_setup_code': row.get('test_setup_code') or '', + }) + logger.info(f'MBPP loaded: {len(samples)} samples') + return samples + + +# ========== Environment lifecycle ========== +def prepare_trajectories( + samples: List[Dict[str, Any]], + pool: ThreadPoolExecutor, +) -> Tuple[List[Dict[str, Any]], List[ToolManager], List[Any]]: + """Create one env per trajectory (reset in parallel) and build trajectories.""" + envs = [backend.make_env() for _ in samples] + # reset() blocks on the network (session handshake, or booting a sandbox); + # run them concurrently so a batch of 32 does not serialize. + list(pool.map(lambda env: env.reset(), envs)) + + trajectories = [] + tool_managers = [] + for sample, env in zip(samples, envs): + tool_managers.append(ToolManager(EnvTool.from_env(env))) + trajectories.append({ + 'messages': [ + {'role': 'system', 'content': backend.SYSTEM_PROMPT}, + {'role': 'user', 'content': sample['prompt']}, + ], + 'tools': backend.TOOL_SCHEMA, + }) + return trajectories, tool_managers, envs + + +def close_envs(envs: List[Any], pool: ThreadPoolExecutor) -> None: + """Release envs so their server-side capacity is freed immediately.""" + list(pool.map(lambda env: env.close(), envs)) + + +def extract_rewards( + envs: List[Any], + samples: List[Dict[str, Any]], + pool: ThreadPoolExecutor, +) -> Tuple[List[float], List[float]]: + """Score submitted solutions against the hidden MBPP tests. + + Shaped so that a full pass dominates, while a submission that passes some + tests still beats one that passes none — which beats never submitting. + + Returns: + ``(rewards, pass_rates)``. + """ + + def score(pair) -> Tuple[float, float]: + env, sample = pair + passed, total = backend.run_tests(env, sample['test_list'], sample['test_setup_code']) + rate = passed / total if total else 0.0 + if total and passed == total: + return 1.0, rate + if getattr(env, 'submitted_code', None): + return 0.1 + 0.4 * rate, rate + return 0.0, rate + + scored = list(pool.map(score, zip(envs, samples))) + return [s[0] for s in scored], [s[1] for s in scored] + + +# swanlab cannot chart string values. TrainMetric / LossMetric emit pre-formatted +# strings ('1.000000e-05', '0.03 iters/s', '30 seconds', '0.8321'); pull the +# leading number out so those keys still render as line charts. Values that do +# not parse are left as-is (they still show up in the text log). +_LEADING_NUMBER_RE = re.compile(r'[-+]?\d*\.?\d+(?:[eE][-+]?\d+)?') + + +def _coerce_for_swanlab(log_dict: Dict[str, Any]) -> Dict[str, Any]: + coerced: Dict[str, Any] = {} + for k, v in log_dict.items(): + if isinstance(v, bool) or isinstance(v, (int, float)): + coerced[k] = v + elif isinstance(v, str) and (m := _LEADING_NUMBER_RE.search(v)): + try: + coerced[k] = float(m.group()) + except ValueError: + coerced[k] = v + else: + coerced[k] = v + return coerced + + +# ========== Main ========== +def main(): + device_groups = [ + DeviceGroup(name='model', ranks=list(range(MODEL_GPUS)), device_type='GPU'), + DeviceGroup(name='sampler', ranks=list(range(MODEL_GPUS, NUM_GPUS)), device_type='GPU'), + ] + model_mesh = DeviceMesh.from_sizes(world_size=MODEL_GPUS, dp_size=MODEL_GPUS) + sampler_mesh = DeviceMesh.from_sizes(world_size=SAMPLER_GPUS, dp_size=SAMPLER_GPUS) + twinkle.initialize(mode='ray', nproc_per_node=NUM_GPUS, groups=device_groups, lazy_collect=False) + + lora_config = LoraConfig( + target_modules='all-linear', + r=LORA_RANK, + lora_alpha=LORA_RANK * 2, + lora_dropout=0.05, + ) + + model = TransformersModel( + model_id=MODEL_ID, + device_mesh=model_mesh, + remote_group='model', + # Mixed-precision training: load master weights in fp32, compute in bf16. + # Without an explicit dtype, from_pretrained falls back to config.dtype + # (bf16 for Qwen), which would keep the master weights in bf16 too — i.e. + # not mixed precision. dtype=fp32 + mixed_precision='bf16' gives fp32 + # weights/optimizer state with bf16 autocast for the forward/backward. + dtype=torch.float32, + mixed_precision='bf16', + ) + model.add_adapter_to_model(ADAPTER_NAME, lora_config, gradient_accumulation_steps=GRADIENT_ACCUMULATION_STEPS) + model.set_optimizer('AdamW', lr=LEARNING_RATE) + model.set_lr_scheduler('CosineAnnealingLR', T_max=MAX_STEPS, eta_min=0) + grpo_epsilon = 0.2 + model.set_loss('GRPOLoss', epsilon=grpo_epsilon) + model.add_metric('GRPOMetric', is_training=True, epsilon=grpo_epsilon) + model.set_processor(InputProcessor, padding_free=True) + model.set_template('Qwen3_5Template', model_id=MODEL_ID, enable_thinking=False) + + sampler = vLLMSampler( + model_id=MODEL_ID, + engine_args={ + 'gpu_memory_utilization': 0.8, + 'max_model_len': 8192, + 'max_lora_rank': 32, + 'enable_lora': True, + 'enable_tower_connector_lora': True, + }, + device_mesh=sampler_mesh, + remote_group='sampler', + ) + sampler.set_template('Qwen3_5Template', model_id=MODEL_ID, enable_thinking=False) + + rollout_template = Qwen3_5Template(MODEL_ID, max_length=8192, enable_thinking=False) + rollout_template.truncation_strategy = 'delete' + + ckpt_manager = CheckpointEngineManager(model=model, sampler=sampler) + + sampling_params = SamplingParams( + max_tokens=MAX_NEW_TOKENS, num_samples=1, logprobs=1, + temperature=1.0, top_p=0.95, + ) + rollout = MultiTurnRollout( + sampler=sampler, + template=rollout_template, + sampling_params=sampling_params, + max_turns=MAX_TURNS, + ) + + advantage_fn = GRPOAdvantage() + metrics = CompletionRewardMetric() + dataset = load_mbpp() + env_pool = ThreadPoolExecutor(max_workers=ENV_CONCURRENCY) + + optim_step = 0 + sample_cursor = 0 + logger.info(f'Starting code RL [{backend.NAME}] — {backend.describe()}') + logger.info(f'Concurrent envs needed: {BATCH_SIZE * NUM_GENERATIONS}') + logger.info(get_device_placement()) + + swanlab.init( + project='twinkle-code', + name=f'code-rl-{backend.NAME}', + config={ + 'backend': backend.NAME, + 'model_id': MODEL_ID, + 'batch_size': BATCH_SIZE, + 'num_generations': NUM_GENERATIONS, + 'lr': LEARNING_RATE, + 'max_steps': MAX_STEPS, + 'lora_rank': LORA_RANK, + 'max_turns': MAX_TURNS, + }, + ) + + while optim_step < MAX_STEPS: + metrics.reset() + + # BATCH_SIZE problems x NUM_GENERATIONS contiguous copies (GRPO groups). + batch = [dataset[(sample_cursor + i) % len(dataset)] for i in range(BATCH_SIZE)] + sample_cursor += BATCH_SIZE + expanded = [s for s in batch for _ in range(NUM_GENERATIONS)] + n_traj = len(expanded) + + # 1. Create envs and build initial trajectories + logger.info(f'[Step {optim_step}] Preparing {n_traj} envs...') + expand_prompts, tool_managers, envs = prepare_trajectories(expanded, env_pool) + + try: + # 2. Sync model weights to sampler + ckpt_manager.sync_weights(merge_and_sync=False) + sampler.reset_prefix_cache() + + # 3. Multi-turn rollout with per-trajectory ToolManagers + all_trajectories: List[Dict[str, Any]] = rollout( + expand_prompts, + tool_manager=tool_managers, + ) + + # 4. Score solutions by replaying the hidden tests in each env + total_rewards, pass_rates = extract_rewards(envs, expanded, env_pool) + finally: + # Envs occupy server-side capacity; always release them. + close_envs(envs, env_pool) + + all_old_logps: List[List[float]] = [] + all_completion_lengths: List[int] = [] + n_turns_per_rollout: List[int] = [] + for traj in all_trajectories: + logprobs = traj.get('logprobs') or [] + all_old_logps.append([lp[0][1] for lp in logprobs] if logprobs else []) + labels = traj.get('labels') or [] + all_completion_lengths.append(sum(1 for l in labels if l != -100)) + n_turns_per_rollout.append(int(traj.get('turns') or 0)) + + # 5. Group-relative advantages + advantages = advantage_fn( + total_rewards, num_generations=NUM_GENERATIONS, scale='group', + ).tolist() + + # 6. Log metrics + metrics.accumulate( + completion_lengths=all_completion_lengths, + rewards={'total': total_rewards}, + ) + avg_reward = sum(total_rewards) / len(total_rewards) if total_rewards else 0.0 + solve_rate = sum(1 for r in total_rewards if r >= 1.0) / max(len(total_rewards), 1) + avg_pass_rate = sum(pass_rates) / len(pass_rates) if pass_rates else 0.0 + avg_turns = sum(n_turns_per_rollout) / len(n_turns_per_rollout) if n_turns_per_rollout else 0.0 + max_turns = max(n_turns_per_rollout) if n_turns_per_rollout else 0 + min_turns = min(n_turns_per_rollout) if n_turns_per_rollout else 0 + logger.info(f'[Step {optim_step}] avg_reward={avg_reward:.3f}, solve_rate={solve_rate:.3f}, ' + f'test_pass_rate={avg_pass_rate:.3f}, avg_turns={avg_turns:.1f}') + + # 7. Filter and train + all_input_data: List[Dict[str, Any]] = [] + filtered_old_logps: List[List[float]] = [] + filtered_advantages: List[float] = [] + max_len = rollout_template.max_length or float('inf') + for i, traj in enumerate(all_trajectories): + traj_len = len(traj.get('input_ids') or traj.get('labels') or []) + comp_len = sum(1 for l in (traj.get('labels') or []) if l != -100) + if traj_len > max_len or comp_len == 0: + continue + all_input_data.append(traj) + filtered_old_logps.append(all_old_logps[i]) + filtered_advantages.append(advantages[i]) + + if len(all_input_data) < MODEL_GPUS: + logger.warning(f'[Step {optim_step}] Only {len(all_input_data)} valid trajectories ' + f'after filtering (need >= {MODEL_GPUS}), skipping this batch.') + continue + + total_completions = len(all_input_data) + logger.info(f'[Step {optim_step}] {total_completions}/{n_traj} trajectories ' + f'passed length filter (max_len={max_len})') + + for mb_start in range(0, total_completions, MINI_BATCH_SIZE): + mb_end = min(mb_start + MINI_BATCH_SIZE, total_completions) + model.forward_backward( + inputs=all_input_data[mb_start:mb_end], + old_logps=filtered_old_logps[mb_start:mb_end], + advantages=filtered_advantages[mb_start:mb_end], + micro_batch_size=MICRO_BATCH_SIZE, + ) + model.clip_grad_and_step() + optim_step += 1 + + if optim_step >= MAX_STEPS: + break + if optim_step % SAVE_STEPS == 0: + model.save(f'code-rl-{backend.NAME}-checkpoint-{optim_step}') + + # 8. Step summary + swanlab logging + log_dict = metrics.calculate() # train/*_reward, train/completion_length + log_dict.update(model.calculate_metric(is_training=True)) # loss, grad_norm, learning rate, token accuracy... + # Code-generation accuracy. solve_rate is pass@1 over the hidden tests + # (fraction of problems passing ALL tests); test_pass_rate is the mean + # per-test pass ratio. Under train/ so swanlab groups them with rewards. + log_dict['train/code_acc'] = solve_rate + log_dict['train/test_pass_rate'] = avg_pass_rate + log_dict['train/avg_reward'] = avg_reward + log_dict['train/avg_turns'] = avg_turns + # Per-rollout turn spread: max_turns hitting MAX_TURNS means rollouts are + # being cut off before submitting; min_turns near 1 means the model often + # submits (or gives up) on the first turn. + log_dict['train/max_turns'] = max_turns + log_dict['train/min_turns'] = min_turns + metrics.reset() + logger.info(f'[Step {optim_step}/{MAX_STEPS}] {log_dict}') + swanlab.log(_coerce_for_swanlab(log_dict), step=optim_step) + + env_pool.shutdown(wait=False) + logger.info(f'Training completed. optim_steps={optim_step}') + model.save(f'code-rl-{backend.NAME}-final') + + +if __name__ == '__main__': + main() diff --git a/cookbook/rl/deploy/README.md b/cookbook/rl/deploy/README.md new file mode 100644 index 000000000..59c2d7810 --- /dev/null +++ b/cookbook/rl/deploy/README.md @@ -0,0 +1,48 @@ +# Deployment templates for Agentic RL + +Network and supervision templates for the three runnable examples in +`cookbook/rl/`. Full reasoning and the selection matrix live in the +[Agentic RL Deployment Guide](../../../docs/source_en/Usage%20Guide/Agentic-RL-Deployment.md). + +## Which example do I run? + +| Need | Example | Environment runs | +|------|---------|------------------| +| Lightweight compute env (games, pure scoring) | `../multi_turn/` | In-process, sharded across Ray workers via `EnvPool` | +| Code execution, env scaled independently of training | `../openenv_code/` | A remote OpenEnv server (HTTP/WebSocket) | +| Real execution semantics (`unittest` + mock, files, `pip`) | `../agentenv/` | Firecracker microVMs via AgentENV | + +Switching examples means changing the adapter class in code. Changing **where** +the environment runs is configuration only — one environment variable +(`OPENENV_BASE_URL` or `AENV_API_URL`). + +## Which network setup? + +Twinkle supports exactly two ways to reach an environment: **direct HTTP** and +**SSH port forwarding**. It ships no VPN / NAT-traversal tooling — that belongs to +your network team's existing infrastructure, and Twinkle only ever sees an +`http://host:port`. + +| Topology | What to do | +|----------|------------| +| Training and env on one host | Bind to localhost: `HOST=127.0.0.1 sh serve.sh`. Nothing else to configure. | +| Separate hosts, same private network / VPC | **The default for most setups.** Bind to the private NIC; security-group inbound allows only the training host. | +| Across networks | Bind to `127.0.0.1` and forward the port over SSH: `ssh -N -L 8000:127.0.0.1:8000 user@env-host`. Point `OPENENV_BASE_URL` at the local entry. | + +## Files + +| File | Purpose | +|------|---------| +| `openenv-server.service` | systemd unit keeping the OpenEnv server alive across crashes and reboots. Install on the **environment** host. | + +## Two things that bite people + +**Neither OpenEnv nor AgentENV has any authentication.** Whoever can reach the +port can execute arbitrary code and consume your capacity. Never bind to +`0.0.0.0` on a reachable interface — restrict it with a security group, or keep +it on `127.0.0.1` and let SSH be the authentication layer. + +**A dropped connection loses the whole rollout batch.** Supervise the server +with `openenv-server.service`, and if you use SSH forwarding, run it under +`autossh -M 0 -N -o ServerAliveInterval=30 ...` rather than in an interactive +shell. diff --git a/cookbook/rl/deploy/openenv-server.service b/cookbook/rl/deploy/openenv-server.service new file mode 100644 index 000000000..5869b0cf8 --- /dev/null +++ b/cookbook/rl/deploy/openenv-server.service @@ -0,0 +1,44 @@ +# OpenEnv environment server, kept alive across reboots and crashes. +# +# A rollout batch dies with the server, so a long training run needs this +# supervised rather than started from an interactive shell. +# +# Install: +# sudo cp openenv-server.service /etc/systemd/system/ +# sudo systemctl daemon-reload +# sudo systemctl enable --now openenv-server +# journalctl -u openenv-server -f +# +# Adjust User, WorkingDirectory and the HOST/WORKERS values before installing. + +[Unit] +Description=OpenEnv code environment server for Twinkle Agentic RL +After=network-online.target +Wants=network-online.target + +[Service] +Type=exec +User=CHANGEME +WorkingDirectory=/path/to/twinkle/cookbook/rl/openenv_code + +# Bind to the private NIC (training on another host in the same network) or +# 127.0.0.1 (same-host training, or reached over SSH port forwarding). +# Never 0.0.0.0 — the server has no authentication. +Environment=HOST=10.0.1.20 +Environment=PORT=8000 +Environment=WORKERS=4 +Environment=MAX_CONCURRENT_ENVS=64 + +ExecStart=/bin/sh serve.sh +Restart=always +RestartSec=5 + +# The server executes model-generated code; keep the blast radius small. +NoNewPrivileges=true +PrivateTmp=true +ProtectSystem=strict +ProtectHome=true +ReadWritePaths=/tmp + +[Install] +WantedBy=multi-user.target diff --git a/cookbook/rl/openenv_code/openenv_code_grpo.py b/cookbook/rl/openenv_code/openenv_code_grpo.py new file mode 100644 index 000000000..0cbc9b555 --- /dev/null +++ b/cookbook/rl/openenv_code/openenv_code_grpo.py @@ -0,0 +1,365 @@ +"""Multi-turn GRPO on MBPP with a remote OpenEnv code-interpreter server. + +Server-mode counterpart of ``cookbook/rl/agentenv/agentenv_grpo.py``: the same +code-writing task, but every trajectory gets a **session on a remote OpenEnv +server** instead of a Firecracker microVM. Compare the two to see what the +execution backend does and does not change. + +Architecture: + - No EnvPool / @remote_class. An OpenEnvClient is a WebSocket client and the + session lives server-side, so Ray sharding buys nothing. The driver creates + one client per trajectory and resets them on a thread pool, because each + connect+reset is a blocking network call. + - One server process hosts many sessions (``MAX_CONCURRENT_ENVS``). Capacity + must cover BATCH_SIZE x NUM_GENERATIONS concurrent trajectories. + +Per-trajectory flow: + 1. reset() opens a session; the server builds a fresh Python interpreter. + 2. MultiTurnRollout drives tool calls: run_python executes code in the + session (state persists across turns); submit_solution records the final + source client-side. + 3. Reward = MBPP unit-test pass rate, measured by replaying the hidden tests + in the same session after the rollout (see ``tools.run_tests``). + 4. GRPO advantages are group-relative across NUM_GENERATIONS rollouts of the + same problem. + +Prerequisites: + 1. Start the environment server (any reachable host, no GPU/KVM needed): + pip install openenv && pip install -e /path/to/OpenEnv/envs/coding_env + sh serve.sh + 2. On the training side: pip install openenv + +Usage: + OPENENV_BASE_URL=http://:8000 sh openenv_code_grpo.sh +""" +import os +from concurrent.futures import ThreadPoolExecutor +from typing import Any, Dict, List, Tuple + +from peft import LoraConfig + +import twinkle +from twinkle import DeviceGroup, DeviceMesh, get_device_placement, get_logger +from twinkle.advantage import GRPOAdvantage +from twinkle.checkpoint_engine import CheckpointEngineManager +from twinkle.cli import CLI +from twinkle.data_format import SamplingParams +from twinkle.metric import CompletionRewardMetric +from twinkle.model import TransformersModel +from twinkle.processor import InputProcessor +from twinkle.sampler import vLLMSampler +from twinkle.template import Qwen3_5Template +from twinkle_agentic.envs import EnvTool, OpenEnvClient +from twinkle_agentic.rollout.multi_turn import MultiTurnRollout +from twinkle_agentic.tools.tool_manager import ToolManager +from tools import SYSTEM_PROMPT, TOOL_SCHEMA, register_tools, run_tests + +logger = get_logger() +args = CLI.from_args() + +# ========== Configuration ========== +MODEL_ID = args.model.model_id or 'ms://Qwen/Qwen3.5-4B' + +MODEL_GPUS = args.infra.model_gpus or 4 +SAMPLER_GPUS = args.infra.sampler_gpus or 4 +NUM_GPUS = MODEL_GPUS + SAMPLER_GPUS + +NUM_GENERATIONS = args.rl.num_generations or 8 +MAX_NEW_TOKENS = args.sampling.max_tokens or 2048 +LEARNING_RATE = args.optimizer.learning_rate or 1e-5 +MAX_STEPS = args.training.max_steps or 1000 +BATCH_SIZE = args.training.batch_size or 4 +MINI_BATCH_SIZE = args.training.mini_batch_size or 8 +MICRO_BATCH_SIZE = args.training.micro_batch_size or 2 +GRADIENT_ACCUMULATION_STEPS = args.training.gradient_accumulation_steps or 1 +ADAPTER_NAME = args.lora.adapter_name or 'default' +SAVE_STEPS = args.training.save_steps or 500 +LORA_RANK = args.lora.lora_r or 16 +MAX_TURNS = int(os.environ.get('MAX_TURNS', '6')) + +# The OpenEnv server. Behind a load balancer this is the LB address; the client +# never needs to know which backend serves its session. +OPENENV_BASE_URL = os.environ.get('OPENENV_BASE_URL', 'http://127.0.0.1:8000') +OPENENV_ENV_NAME = os.environ.get('OPENENV_ENV_NAME', 'coding_env') +# Parallelism for the blocking connect/reset/score calls in the driver. +ENV_CONCURRENCY = int(os.environ.get('ENV_CONCURRENCY', '16')) + + +# ========== Dataset (MBPP) ========== +def load_mbpp() -> List[Dict[str, Any]]: + """Load MBPP as [{'prompt', 'test_list', 'test_setup_code'}].""" + from modelscope.msdatasets import MsDataset + raw = MsDataset.load('opencompass/mbpp', subset_name='full', split='train') + samples = [] + for row in raw: + test_list = list(row['test_list'] or []) + if not test_list: + continue + # The first assertion is handed to the model as a signature hint: MBPP + # descriptions alone do not pin the function name, and the hidden tests + # call it by name. + prompt = (f"{row['text']}\n\n" + f'Your function must satisfy this call:\n{test_list[0]}') + samples.append({ + 'prompt': prompt, + 'test_list': test_list, + 'test_setup_code': row.get('test_setup_code') or '', + }) + logger.info(f'MBPP loaded: {len(samples)} samples') + return samples + + +# ========== Environment Setup ========== +def make_env() -> OpenEnvClient: + """One session per trajectory; only the two task tools are exposed.""" + env = OpenEnvClient( + env_name=OPENENV_ENV_NAME, + base_url=OPENENV_BASE_URL, + tools=[TOOL_SCHEMA[0]], # run_python; submit_solution is registered below + # Code execution can be slow; keep the per-message timeout generous. + message_timeout_s=120.0, + ) + return register_tools(env) + + +def prepare_trajectories( + samples: List[Dict[str, Any]], + pool: ThreadPoolExecutor, +) -> Tuple[List[Dict[str, Any]], List[ToolManager], List[OpenEnvClient]]: + """Open one session per trajectory (in parallel) and build trajectories.""" + envs = [make_env() for _ in samples] + # reset() blocks on the WebSocket handshake plus server-side env creation; + # run them concurrently so a batch of 32 does not serialize. + list(pool.map(lambda env: env.reset(), envs)) + + trajectories = [] + tool_managers = [] + for sample, env in zip(samples, envs): + tool_managers.append(ToolManager(EnvTool.from_env(env))) + trajectories.append({ + 'messages': [ + {'role': 'system', 'content': SYSTEM_PROMPT}, + {'role': 'user', 'content': sample['prompt']}, + ], + 'tools': TOOL_SCHEMA, + }) + return trajectories, tool_managers, envs + + +def close_envs(envs: List[OpenEnvClient], pool: ThreadPoolExecutor) -> None: + """Close all sessions so their server-side slots are freed immediately.""" + list(pool.map(lambda env: env.close(), envs)) + + +def extract_rewards( + envs: List[OpenEnvClient], + samples: List[Dict[str, Any]], + pool: ThreadPoolExecutor, +) -> Tuple[List[float], List[float]]: + """Score submitted solutions against the hidden MBPP tests. + + Shaped so that a full pass dominates, while a submission that passes some + tests still beats one that passes none — which beats never submitting. + + Returns: + ``(rewards, pass_rates)``. + """ + + def score(pair) -> Tuple[float, float]: + env, sample = pair + passed, total = run_tests(env, sample['test_list'], sample['test_setup_code']) + rate = passed / total if total else 0.0 + if total and passed == total: + return 1.0, rate + if getattr(env, 'submitted_code', None): + return 0.1 + 0.4 * rate, rate + return 0.0, rate + + scored = list(pool.map(score, zip(envs, samples))) + return [s[0] for s in scored], [s[1] for s in scored] + + +# ========== Main ========== +def main(): + device_groups = [ + DeviceGroup(name='model', ranks=list(range(MODEL_GPUS)), device_type='GPU'), + DeviceGroup(name='sampler', ranks=list(range(MODEL_GPUS, NUM_GPUS)), device_type='GPU'), + ] + model_mesh = DeviceMesh.from_sizes(world_size=MODEL_GPUS, dp_size=MODEL_GPUS) + sampler_mesh = DeviceMesh.from_sizes(world_size=SAMPLER_GPUS, dp_size=SAMPLER_GPUS) + twinkle.initialize(mode='ray', nproc_per_node=NUM_GPUS, groups=device_groups, lazy_collect=False) + + lora_config = LoraConfig( + target_modules='all-linear', + r=LORA_RANK, + lora_alpha=LORA_RANK * 2, + lora_dropout=0.05, + ) + + model = TransformersModel( + model_id=MODEL_ID, + device_mesh=model_mesh, + remote_group='model', + ) + model.add_adapter_to_model(ADAPTER_NAME, lora_config, gradient_accumulation_steps=GRADIENT_ACCUMULATION_STEPS) + model.set_optimizer('AdamW', lr=LEARNING_RATE) + model.set_lr_scheduler('CosineAnnealingLR', T_max=MAX_STEPS, eta_min=0) + model.set_loss('GRPOLoss', epsilon=0.2) + model.set_processor(InputProcessor, padding_free=True) + model.set_template('Qwen3_5Template', model_id=MODEL_ID, enable_thinking=False) + + sampler = vLLMSampler( + model_id=MODEL_ID, + engine_args={ + 'gpu_memory_utilization': 0.8, + 'max_model_len': 8192, + 'max_lora_rank': 32, + 'enable_lora': True, + 'enable_tower_connector_lora': True, + }, + device_mesh=sampler_mesh, + remote_group='sampler', + ) + sampler.set_template('Qwen3_5Template', model_id=MODEL_ID, enable_thinking=False) + + rollout_template = Qwen3_5Template(MODEL_ID, max_length=8192, enable_thinking=False) + rollout_template.truncation_strategy = 'delete' + + ckpt_manager = CheckpointEngineManager(model=model, sampler=sampler) + + sampling_params = SamplingParams( + max_tokens=MAX_NEW_TOKENS, num_samples=1, logprobs=1, + temperature=1.0, top_p=0.95, + ) + rollout = MultiTurnRollout( + sampler=sampler, + template=rollout_template, + sampling_params=sampling_params, + max_turns=MAX_TURNS, + ) + + advantage_fn = GRPOAdvantage() + metrics = CompletionRewardMetric() + dataset = load_mbpp() + env_pool = ThreadPoolExecutor(max_workers=ENV_CONCURRENCY) + + optim_step = 0 + sample_cursor = 0 + logger.info(f'Starting OpenEnv (server mode) GRPO: base_url={OPENENV_BASE_URL}, env={OPENENV_ENV_NAME}') + logger.info(f'Concurrent sessions needed: {BATCH_SIZE * NUM_GENERATIONS} ' + f'(server capacity = workers x MAX_CONCURRENT_ENVS)') + logger.info(get_device_placement()) + + while optim_step < MAX_STEPS: + metrics.reset() + + # BATCH_SIZE problems x NUM_GENERATIONS contiguous copies (GRPO groups). + batch = [dataset[(sample_cursor + i) % len(dataset)] for i in range(BATCH_SIZE)] + sample_cursor += BATCH_SIZE + expanded = [s for s in batch for _ in range(NUM_GENERATIONS)] + n_traj = len(expanded) + + # 1. Open sessions and build initial trajectories + logger.info(f'[Step {optim_step}] Opening {n_traj} sessions...') + expand_prompts, tool_managers, envs = prepare_trajectories(expanded, env_pool) + + try: + # 2. Sync model weights to sampler + ckpt_manager.sync_weights(merge_and_sync=False) + sampler.reset_prefix_cache() + + # 3. Multi-turn rollout with per-trajectory ToolManagers + all_trajectories: List[Dict[str, Any]] = rollout( + expand_prompts, + tool_manager=tool_managers, + ) + + # 4. Score solutions by replaying the hidden tests in each session + total_rewards, pass_rates = extract_rewards(envs, expanded, env_pool) + finally: + # Sessions occupy server capacity; always release them. + close_envs(envs, env_pool) + + all_old_logps: List[List[float]] = [] + all_completion_lengths: List[int] = [] + n_turns_per_rollout: List[int] = [] + for traj in all_trajectories: + logprobs = traj.get('logprobs') or [] + all_old_logps.append([lp[0][1] for lp in logprobs] if logprobs else []) + labels = traj.get('labels') or [] + all_completion_lengths.append(sum(1 for l in labels if l != -100)) + n_turns_per_rollout.append(int(traj.get('turns') or 0)) + + # 5. Group-relative advantages + advantages = advantage_fn( + total_rewards, num_generations=NUM_GENERATIONS, scale='group', + ).tolist() + + # 6. Log metrics + metrics.accumulate( + completion_lengths=all_completion_lengths, + rewards={'total': total_rewards}, + ) + avg_reward = sum(total_rewards) / len(total_rewards) if total_rewards else 0.0 + solve_rate = sum(1 for r in total_rewards if r >= 1.0) / max(len(total_rewards), 1) + avg_pass_rate = sum(pass_rates) / len(pass_rates) if pass_rates else 0.0 + avg_turns = sum(n_turns_per_rollout) / len(n_turns_per_rollout) if n_turns_per_rollout else 0.0 + logger.info(f'[Step {optim_step}] avg_reward={avg_reward:.3f}, solve_rate={solve_rate:.3f}, ' + f'test_pass_rate={avg_pass_rate:.3f}, avg_turns={avg_turns:.1f}') + + # 7. Filter and train (same recipe as the AgentENV example) + all_input_data: List[Dict[str, Any]] = [] + filtered_old_logps: List[List[float]] = [] + filtered_advantages: List[float] = [] + max_len = rollout_template.max_length or float('inf') + for i, traj in enumerate(all_trajectories): + traj_len = len(traj.get('input_ids') or traj.get('labels') or []) + comp_len = sum(1 for l in (traj.get('labels') or []) if l != -100) + if traj_len > max_len or comp_len == 0: + continue + all_input_data.append(traj) + filtered_old_logps.append(all_old_logps[i]) + filtered_advantages.append(advantages[i]) + + if len(all_input_data) < MODEL_GPUS: + logger.warning(f'[Step {optim_step}] Only {len(all_input_data)} valid trajectories ' + f'after filtering (need >= {MODEL_GPUS}), skipping this batch.') + continue + + total_completions = len(all_input_data) + logger.info(f'[Step {optim_step}] {total_completions}/{n_traj} trajectories ' + f'passed length filter (max_len={max_len})') + + for mb_start in range(0, total_completions, MINI_BATCH_SIZE): + mb_end = min(mb_start + MINI_BATCH_SIZE, total_completions) + model.forward_backward( + inputs=all_input_data[mb_start:mb_end], + old_logps=filtered_old_logps[mb_start:mb_end], + advantages=filtered_advantages[mb_start:mb_end], + micro_batch_size=MICRO_BATCH_SIZE, + ) + model.clip_grad_and_step() + optim_step += 1 + + if optim_step >= MAX_STEPS: + break + if optim_step % SAVE_STEPS == 0: + model.save(f'openenv-code-grpo-checkpoint-{optim_step}') + + # 8. Step summary + log_dict = metrics.calculate() + log_dict.update(model.calculate_metric(is_training=True)) + log_dict['avg_turns'] = avg_turns + log_dict['avg_reward'] = avg_reward + log_dict['solve_rate'] = solve_rate + log_dict['test_pass_rate'] = avg_pass_rate + metrics.reset() + logger.info(f'[Step {optim_step}/{MAX_STEPS}] {log_dict}') + + env_pool.shutdown(wait=False) + logger.info(f'Training completed. optim_steps={optim_step}') + model.save('openenv-code-grpo-final') + + +if __name__ == '__main__': + main() diff --git a/cookbook/rl/openenv_code/openenv_code_grpo.sh b/cookbook/rl/openenv_code/openenv_code_grpo.sh new file mode 100644 index 000000000..62d3d9221 --- /dev/null +++ b/cookbook/rl/openenv_code/openenv_code_grpo.sh @@ -0,0 +1,48 @@ +#!/bin/sh +set -eu + +# Multi-turn GRPO on MBPP with a remote OpenEnv code-interpreter server. +# +# The model writes a Python function, tests it in a remote interpreter session, +# and submits it; reward is the hidden unit-test pass rate. The environment runs +# as an ordinary HTTP/WebSocket service — no Docker, no KVM, no GPU on that side. +# +# One-time setup: +# 1. On the environment host: +# pip install openenv +# pip install -e /path/to/OpenEnv/envs/coding_env +# sh serve.sh # note the URL it serves on +# 2. On the training host: +# pip install openenv +# +# Capacity check: the server must host BATCH_SIZE x NUM_GENERATIONS concurrent +# sessions (32 with the defaults below). serve.sh gives WORKERS x +# MAX_CONCURRENT_ENVS = 256 by default, so there is headroom. +# +# Run (override anything at invocation): +# OPENENV_BASE_URL=http://10.0.0.5:8000 sh openenv_code_grpo.sh --max-steps 500 + +# ---- OpenEnv server (a load-balancer address works here too) ---- +export OPENENV_BASE_URL="${OPENENV_BASE_URL:-http://127.0.0.1:8000}" +# Environment package providing the client + Action classes. +export OPENENV_ENV_NAME="${OPENENV_ENV_NAME:-coding_env}" +# Concurrent connect/reset/score calls from the driver. +export ENV_CONCURRENCY="${ENV_CONCURRENCY:-16}" +# Max tool-calling turns per episode. +export MAX_TURNS="${MAX_TURNS:-6}" + +python openenv_code_grpo.py \ + --model-id ms://Qwen/Qwen3.5-4B \ + --model-gpus 4 \ + --sampler-gpus 4 \ + --num-generations 8 \ + --max-tokens 2048 \ + --batch-size 4 \ + --mini-batch-size 8 \ + --micro-batch-size 2 \ + --max-steps 1000 \ + --lr 1e-5 \ + --lora-r 16 \ + --save-steps 500 \ + --adapter-name default \ + "$@" diff --git a/cookbook/rl/openenv_code/serve.sh b/cookbook/rl/openenv_code/serve.sh new file mode 100644 index 000000000..ca0c4db1c --- /dev/null +++ b/cookbook/rl/openenv_code/serve.sh @@ -0,0 +1,52 @@ +#!/bin/sh +set -eu + +# Start the OpenEnv code environment server for cookbook/rl/openenv_code. +# +# This is the REMOTE side: run it on any machine reachable from the training +# driver — no Docker, no KVM, no GPU. The training script only needs the +# resulting URL. +# +# One-time setup: +# pip install openenv +# pip install -e /path/to/OpenEnv/envs/coding_env +# +# Usage: +# sh serve.sh # 4 workers x 64 sessions = 256 sessions +# WORKERS=8 PORT=9000 sh serve.sh +# HOST=127.0.0.1 sh serve.sh # same-machine training: no network exposure +# +# SECURITY: OpenEnv has no authentication. Anyone who can reach this port can +# execute code in the sandbox and consume your capacity. Either bind to +# 127.0.0.1 (when training runs on this same host), or restrict the port to the +# training host's IP with a security group / firewall rule. Never leave the +# default 0.0.0.0 reachable from the public internet. +# +# Then point the training script at it: +# OPENENV_BASE_URL=http://:8000 sh openenv_code_grpo.sh + +HOST="${HOST:-0.0.0.0}" +PORT="${PORT:-8000}" +# Uvicorn worker processes. Each worker holds its own app instance, so total +# capacity is WORKERS x MAX_CONCURRENT_ENVS. Size it above the number of +# concurrent trajectories (BATCH_SIZE x NUM_GENERATIONS), otherwise the extra +# WebSocket connections are rejected at capacity. +WORKERS="${WORKERS:-4}" +export MAX_CONCURRENT_ENVS="${MAX_CONCURRENT_ENVS:-64}" + +cd "$(dirname "$0")" + +echo "Serving twinkle_code_env on ${HOST}:${PORT}" +echo " workers=${WORKERS}, max_concurrent_envs=${MAX_CONCURRENT_ENVS} per worker" +echo " capacity=$((WORKERS * MAX_CONCURRENT_ENVS)) concurrent sessions" +if [ "$HOST" = "0.0.0.0" ]; then + echo " WARNING: bound to 0.0.0.0 with no authentication — restrict this port" + echo " to the training host, or use HOST=127.0.0.1 if training is local." +fi + +exec uvicorn server_app:app --host "$HOST" --port "$PORT" --workers "$WORKERS" + +# Docker alternative (build coding_env's image, then override the app module): +# docker build -t coding-env:latest -f /path/to/OpenEnv/envs/coding_env/server/Dockerfile /path/to/OpenEnv +# Note the image serves upstream coding_env.server.app, which is capped at one +# session; mount this file and point uvicorn at it to keep the concurrency fix. diff --git a/cookbook/rl/openenv_code/server_app.py b/cookbook/rl/openenv_code/server_app.py new file mode 100644 index 000000000..2b2a000d7 --- /dev/null +++ b/cookbook/rl/openenv_code/server_app.py @@ -0,0 +1,113 @@ +"""OpenEnv server for the code-writing RL task — server mode, no Docker needed. + +Wraps OpenEnv's ``PythonCodeActEnv`` (``OpenEnv/envs/coding_env``) with the +three changes a training workload needs. Each one is a deliberate deviation +from upstream defaults: + +1. ``SUPPORTS_CONCURRENT_SESSIONS = True``. Upstream leaves this at the + conservative default, which caps the server at ONE session + (``create_app`` raises ``ConcurrencyConfigurationError`` for + ``max_concurrent_envs > 1`` otherwise). The class is in fact session-safe: + ``__init__`` builds a private executor and state and shares nothing, and + ``create_app`` receives the class as a *factory*, so every WebSocket + connection gets a fresh instance. + +2. A wider import whitelist. smolagents' ``LocalPythonExecutor`` authorises + only ``json`` by default, so ``import math`` / ``collections`` — which many + MBPP solutions need — would fail. + +3. No reward transforms. ``coding_env.create_safe_coding_transform()`` + overwrites ``observation.reward`` with code-style heuristics (-1.0 when the + code matches ``open(`` / ``import os``, +0.1 for short code). For this task + the reward must come from unit tests, computed by the training script, so a + style score on the same channel is noise. + +Prerequisites:: + + pip install openenv + pip install -e /path/to/OpenEnv/envs/coding_env # brings in smolagents + +Run:: + + sh serve.sh + # or explicitly: + MAX_CONCURRENT_ENVS=64 uvicorn server_app:app --host 0.0.0.0 --port 8000 --workers 4 + +Environment variables: + MAX_CONCURRENT_ENVS: concurrent sessions per worker process (default 64). +""" +import os + +from coding_env.models import CodeAction, CodeObservation +from coding_env.server.python_codeact_env import PythonCodeActEnv +from coding_env.server.python_executor import PyExecutor +from openenv.core.env_server import create_app + +# Modules the sandboxed executor may import. Keep this list tight: it is the +# only thing standing between model-generated code and this process, since +# LocalPythonExecutor is an AST interpreter, not an OS-level sandbox. +ALLOWED_IMPORTS = [ + 'math', + 're', + 'collections', + 'itertools', + 'functools', + 'operator', + 'heapq', + 'bisect', + 'string', + 'statistics', + 'fractions', + 'decimal', + 'datetime', + 'copy', + 'json', +] + +MAX_CONCURRENT_ENVS = int(os.environ.get('MAX_CONCURRENT_ENVS', '64')) + + +class ConcurrentCodeEnv(PythonCodeActEnv): + """Session-isolated Python executor with a task-appropriate config.""" + + SUPPORTS_CONCURRENT_SESSIONS = True + + def __init__(self): + super().__init__() + self._configure() + + def reset(self, **kwargs): + # The parent's reset() rebuilds the executor and the transform with + # upstream defaults, so re-apply our config afterwards. **kwargs + # absorbs the seed / episode_id the server forwards from the client. + observation = super().reset() + self._configure() + return observation + + def _configure(self) -> None: + self._executor = PyExecutor(additional_imports=list(ALLOWED_IMPORTS)) + # Drop the style/safety reward heuristics; reward comes from the tests. + self.transform = None + + +app = create_app( + ConcurrentCodeEnv, + CodeAction, + CodeObservation, + env_name='twinkle_code_env', + max_concurrent_envs=MAX_CONCURRENT_ENVS, +) + + +def main(): + """Entry point for ``python server_app.py`` (single worker).""" + import uvicorn + uvicorn.run( + app, + host=os.environ.get('HOST', '0.0.0.0'), + port=int(os.environ.get('PORT', '8000')), + ) + + +if __name__ == '__main__': + main() diff --git a/cookbook/rl/openenv_code/tools.py b/cookbook/rl/openenv_code/tools.py new file mode 100644 index 000000000..b3df9c48e --- /dev/null +++ b/cookbook/rl/openenv_code/tools.py @@ -0,0 +1,138 @@ +"""Tool definitions for OpenEnv-backed code-writing RL (MBPP). + +Same role as ``cookbook/rl/agentenv/tools.py``, but the tools run against a +remote OpenEnv **server session** instead of a Firecracker sandbox. The pair of +examples is intentionally symmetric: identical task, identical tool names, +different execution backend. + +Two tools are exposed to the model: + * ``run_python`` — execute a snippet in the session (server-backed). + * ``submit_solution`` — hand in the final function (recorded client-side). + +The session namespace persists across ``run_python`` calls, so the model can +define a function in one turn and probe it in the next. +""" +from typing import Any, Dict, List, Tuple + +from twinkle_agentic.envs import OpenEnvClient + +SYSTEM_PROMPT = """You are an expert Python programmer with access to a Python interpreter. + +Write a function that solves the given task, verify it, then submit it. + +Rules: +- Use `run_python` to define and test your function. The interpreter keeps its + state between calls, so a function defined in one call stays available. +- Available modules: math, re, collections, itertools, functools, operator, + heapq, bisect, string, statistics, fractions, decimal, datetime, copy, json. + There is no file or network access. +- Match the function name and signature implied by the task description and + the example call, otherwise the hidden tests cannot find your function. +- When your function works, call `submit_solution` with the COMPLETE final + source (imports plus the function definition). +- After submitting, reply with one short sentence and do NOT call any more tools. +- You have a limited number of turns, so do not run redundant code.""" + +TOOL_SCHEMA: List[Dict[str, Any]] = [ + { + 'type': 'function', + 'function': { + 'name': 'run_python', + 'description': 'Execute a Python snippet in the interpreter and return its stdout/stderr.', + 'parameters': { + 'type': 'object', + 'properties': { + 'code': { + 'type': 'string', + 'description': 'Python source to execute. Use print(...) to inspect values.', + }, + }, + 'required': ['code'], + }, + }, + }, + { + 'type': 'function', + 'function': { + 'name': 'submit_solution', + 'description': 'Submit the final solution source and end the coding phase.', + 'parameters': { + 'type': 'object', + 'properties': { + 'code': { + 'type': 'string', + 'description': 'Complete final source: imports plus the function definition.', + }, + }, + 'required': ['code'], + }, + }, + }, +] + + +def _submit_solution(env: OpenEnvClient, arguments: Dict[str, Any]) -> str: + """Record the solution on the env; the training loop scores it later.""" + code = (arguments.get('code') or '').strip() + if not code: + return "Error: 'code' argument is required." + env.submitted_code = code + return 'Solution submitted.' + + +def register_tools(env: OpenEnvClient) -> OpenEnvClient: + """Attach the task tools to a fresh OpenEnvClient instance. + + ``run_python`` stays server-backed: the default action mapper turns + ``{'code': ...}`` straight into the env's ``CodeAction``. + """ + env.submitted_code = None + return env.register_tool(TOOL_SCHEMA[1], _submit_solution) + + +def _last_line(text: str) -> str: + lines = [line for line in (text or '').splitlines() if line.strip()] + return lines[-1].strip() if lines else '' + + +def _ok(result) -> bool: + return not (getattr(result.observation, 'exit_code', 0) or 0) + + +def run_tests(env: OpenEnvClient, test_list: List[str], setup_code: str = '') -> Tuple[int, int]: + """Run the hidden tests against the submitted solution in the same session. + + Each assertion is executed as its own ``print()`` step rather than as + one block of ``assert`` statements. Two reasons: printing the boolean tells + "the assertion evaluated to False" apart from "the code blew up", whereas a + failed ``assert`` surfaces as an exception indistinguishable from a crash + inside the solution; and one step per test isolates a test that raises, so + the remaining ones still run. + + The executor does support ``assert``/``try`` (smolagents implements both), + so this is a diagnosability choice, not a capability workaround. + + Returns: + ``(n_passed, n_total)``. ``(0, n)`` when nothing was submitted or the + solution itself fails to execute. + """ + total = len(test_list) + solution = getattr(env, 'submitted_code', None) + if not solution: + return 0, total + + # Define the submitted solution in the session namespace. + if not _ok(env.execute({'code': solution})): + return 0, total + if setup_code and not _ok(env.execute({'code': setup_code})): + return 0, total + + passed = 0 + for test in test_list: + expr = test.strip() + if expr.startswith('assert '): + expr = expr[len('assert '):] + result = env.execute({'code': f'print({expr})'}) + if _ok(result) and _last_line(getattr(result.observation, 'stdout', '')) == 'True': + passed += 1 + return passed, total diff --git a/docs/source_en/Components/Agentic/Envs.md b/docs/source_en/Components/Agentic/Envs.md index 3e2e90b3a..4a47b001f 100644 --- a/docs/source_en/Components/Agentic/Envs.md +++ b/docs/source_en/Components/Agentic/Envs.md @@ -97,41 +97,100 @@ env_tool = EnvTool( ) ``` -## OpenEnv +## OpenEnv: Two Integration Modes -`OpenEnv` adapts an [OpenEnv](https://github.com/OpenEnv) WebSocket-based environment server as a synchronous Twinkle `Env`. +An [OpenEnv](https://github.com/meta-pytorch/OpenEnv) environment package ships both an `Environment` implementation and an `EnvClient`, so Twinkle provides two adapters for two very different deployment shapes: + +| | **Embedded** `OpenEnv` | **Server** `OpenEnvClient` | +|---|---|---| +| Where the env runs | Inside the training process (Environment instantiated directly) | A separate OpenEnv service process / container | +| Transport | None (local function calls) | Persistent WebSocket (one connection = one session) | +| Isolation | None; shares memory and the GPU node with training | Process / container level; can live on a completely separate machine | +| Dependencies | Env package must be installed on the training node | Env package only needed on the environment node | +| Scaling | `EnvPool` shards across Ray workers | The server's own concurrent sessions + replicas | +| Best for | Pure-compute lightweight envs (board games, text games) | Code execution, anything needing isolation or independent scaling | + +> Do **not** wrap `OpenEnvClient` in `EnvPool`: the session's lifetime is owned by the server, so sharding it again through Ray buys nothing and only adds an RPC hop. + +### Mode 1: Embedded `OpenEnv` + +Bypasses OpenEnv's FastAPI server and constructs the Environment in-process — zero network overhead. ```python from twinkle_agentic.envs.openenv import OpenEnv env = OpenEnv( - base_url='http://localhost:8000', - env_cls='coding_env.CodingEnv', # Optional typed client - env_kwargs={'message_timeout_s': 30}, - tool_schema=[...], # Optional tool definitions + env_name='openspiel_env', # Package name; Environment / Action classes auto-discovered + env_kwargs={'game_name': 'blackjack'}, # Passed to the Environment constructor ) +result = env.reset() +result = env.step('play', {'action': 'hit'}) ``` -### Parameters +| Parameter | Type | Description | +|-----------|------|-------------| +| `env_name` | `str` | OpenEnv package name (e.g. `'coding_env'`). The `*Environment` class is discovered from `.server` and the `*Action` class from the package's `__all__`. | +| `env_cls` | `str` or class | Explicit Environment class (`'module:ClassName'`), used instead of `env_name`. | +| `env_kwargs` | `Dict` | Kwargs for the Environment constructor. | +| `action_cls` | `str` or class | Explicit Action class; auto-discovered from `env_name` when omitted. | +| `action_mapper` | `Callable` | `(tool_name, arguments) -> action`. Defaults to passing the tool arguments as Action fields. | + +### Mode 2: Server `OpenEnvClient` + +First, run the OpenEnv environment as an ordinary HTTP/WebSocket service on the environment host (no Docker required): + +```bash +pip install openenv +pip install -e /path/to/OpenEnv/envs/coding_env +uvicorn coding_env.server.app:app --host 0.0.0.0 --port 8000 --workers 4 +``` + +Then create one client per trajectory on the training side. Each instance owns its own WebSocket session, and the server keeps a dedicated Environment instance for it: + +```python +from twinkle_agentic.envs.openenv import OpenEnvClient + +env = OpenEnvClient( + env_name='coding_env', # EnvClient subclass + Action class auto-discovered + base_url='http://10.0.0.5:8000', # Or set OPENENV_BASE_URL + message_timeout_s=120, # Raise it when the env runs long-executing code +) +env.reset() +result = env.step('run_python', {'code': 'print(1 + 1)'}) +print(result.observation) # '2' +env.close() +``` | Parameter | Type | Description | |-----------|------|-------------| -| `base_url` | `str` | URL of the running OpenEnv server. | -| `env_cls` | `str` or class | Dotted import path or class for a typed client. `None` uses `GenericEnvClient`. | -| `env_kwargs` | `Dict` | Extra kwargs for the client constructor. | -| `tool_schema` | `List[ToolInfo]` | Tool definitions exposed via `tools()`. | -| `action_mapper` | `Callable` | Custom function to map `(tool_name, args)` to the action dict sent to the server. | +| `env_name` | `str` | OpenEnv package name; the `EnvClient` subclass and `*Action` class are auto-discovered from it. | +| `env_cls` | `str` or class | Explicit client class (`'module:ClassName'`), used instead of `env_name`. | +| `base_url` | `str` | Server address, `http(s)://` or `ws(s)://` (converted automatically). Falls back to `OPENENV_BASE_URL`. A load-balancer address works here too. | +| `action_cls` | `str` or class | Action class; auto-discovered when omitted. | +| `action_mapper` | `Callable` | `(tool_name, arguments) -> action`, returning an Action instance or a dict of fields. | +| `tools` | `List[ToolInfo]` | Tool schemas exposed to the model. Defaults to a single `run_python(code)`, matching OpenEnv's code environments. | +| `reset_kwargs` | `Dict` | Kwargs forwarded to the server's `reset()` (e.g. `task_prompt` / `expected_answer` for `repl_env`). Also settable per episode via the `env.reset_kwargs` attribute. | +| `connect_timeout_s` | `float` | WebSocket connect timeout, default 10s. | +| `message_timeout_s` | `float` | Per-message timeout, default 120s. | +| `client_kwargs` | `Dict` | Extra kwargs for the OpenEnv client constructor. | + +Additional capabilities: + +- `register_tool(tool_info, handler)`: register a tool handled **locally on the client** instead of being sent to the server. The typical use is a bookkeeping tool such as `submit_solution`, which records the model's answer on the env for the training loop to score later. A tool of the same name shadows the default one. +- `execute(action)`: send an action and return the server's **raw** `StepResult`, so you can read typed fields such as `exit_code`. The training loop uses it to run scoring code inside the same session. +- `episode_reward` / `last_result` / `client`: cumulative reward, last raw result, and the underlying synchronous client. + +**Capacity and concurrency**: `OpenEnvClient` calls the OpenEnv client's `.sync()`, giving every instance a dedicated background event loop, so a thread pool of concurrent resets/steps is safe. But the server must have room for every concurrent session: the environment class has to declare `SUPPORTS_CONCURRENT_SESSIONS = True` and `create_app(..., max_concurrent_envs=N)` must be large enough, otherwise extra connections are rejected. OpenEnv's bundled `coding_env` is **single-session** by default (`SUPPORTS_CONCURRENT_SESSIONS = False`) and needs to be subclassed to lift that; `cookbook/rl/openenv_code/server_app.py` shows the full pattern. ### Usage with Rollout +Downstream usage is identical for both modes: + ```python -from twinkle_agentic.envs.openenv import OpenEnv from twinkle_agentic.envs.env_tool import EnvTool from twinkle_agentic.tools.tool_manager import ToolManager from twinkle_agentic.rollout.api_multi_turn import APIMultiTurnRollout -# Set up environment -env = OpenEnv(base_url='http://localhost:8000', tool_schema=[...]) env.reset() # Bridge to ToolManager @@ -143,7 +202,36 @@ rollout = APIMultiTurnRollout(api=api, tool_manager=manager, max_turns=10) results = rollout(trajectories) ``` -### Implementing a Custom Environment +See [Agentic RL Best Practices](../../Usage%20Guide/Agentic-RL-Best-Practices.md) for an end-to-end multi-turn GRPO example. + +## EnvPool: Distributed Environment Pool + +`EnvPool` is a `@remote_class` that shards `pool_size` **embedded** `OpenEnv` instances across Ray workers. Each worker manages `pool_size // world_size` slots, and `reset` / `step` are routed to the owning worker automatically via `remote_function`. + +```python +from twinkle_agentic.envs.openenv import EnvPool + +pool = EnvPool( + pool_size=64, + device_mesh=mesh, + env_kwargs={'env_name': 'openspiel_env', 'env_kwargs': {'game_name': 'blackjack'}}, +) + +# Each slot becomes a standard Env, ready for EnvTool / ToolManager +envs = pool.get_adapters(64) +env_tools = EnvTool.from_env(envs[0]) +``` + +| Method | Description | +|--------|-------------| +| `reset(idx)` / `step(idx, tool_name, arguments)` | Operate on a single slot. | +| `reset_batch(indices)` / `step_batch(indices, tool_names, arguments_list)` | Batch operations covering many slots in one RPC, returned in `indices` order. | +| `get_adapters(n)` | Wrap the first `n` slots as `EnvPoolAdapter` (a standard `Env`) on the driver side. | +| `close()` | Close all environments. | + +`EnvPoolAdapter` implements the standard `Env` interface and proxies `reset` / `step` to the owning worker. On a `step` failure it returns `done=True` with the error in `info['error']`, so one broken environment does not stall the whole rollout batch. + +## Implementing a Custom Environment ```python from twinkle_agentic.envs.base import Env, StepResult diff --git a/docs/source_en/Components/Agentic/Multi-Turn-Tool-Usage.md b/docs/source_en/Components/Agentic/Multi-Turn-Tool-Usage.md index 06a962a33..24296f584 100644 --- a/docs/source_en/Components/Agentic/Multi-Turn-Tool-Usage.md +++ b/docs/source_en/Components/Agentic/Multi-Turn-Tool-Usage.md @@ -143,25 +143,26 @@ rollout = APIMultiTurnRollout(api=api, tool_manager=manager, max_turns=10) ## Using OpenEnv Environments -Connect to a remote OpenEnv WebSocket server: +Connect to a running OpenEnv server (each client instance maps to one server-side session): ```python -from twinkle_agentic.envs.openenv import OpenEnv +from twinkle_agentic.envs.openenv import OpenEnvClient from twinkle_agentic.envs.env_tool import EnvTool -env = OpenEnv( +env = OpenEnvClient( + env_name='coding_env', base_url='http://localhost:8000', - env_cls='coding_env.CodingEnv', - tool_schema=[{ + tools=[{ 'type': 'function', 'function': { - 'name': 'submit', - 'description': 'Submit code solution.', + 'name': 'run_python', + 'description': 'Execute Python code in the remote interpreter.', 'parameters': { 'type': 'object', 'properties': { 'code': {'type': 'string'}, }, + 'required': ['code'], }, }, }], @@ -172,6 +173,8 @@ env_tools = EnvTool.from_env(env) manager = ToolManager(env_tools) ``` +To run the environment inside the training process instead (no isolation, zero network overhead), swap `OpenEnvClient` for `OpenEnv(env_name=..., env_kwargs=...)`. See [Environments](./Envs.md) for a full comparison of the two modes. + ## Per-Trajectory Tool Managers For scenarios where each trajectory needs its own tool set (e.g., trajectory-bound state): diff --git a/docs/source_en/Usage Guide/Agentic-RL-Best-Practices.md b/docs/source_en/Usage Guide/Agentic-RL-Best-Practices.md new file mode 100644 index 000000000..2472addc5 --- /dev/null +++ b/docs/source_en/Usage Guide/Agentic-RL-Best-Practices.md @@ -0,0 +1,305 @@ +# Agentic RL Best Practices + +This guide walks through a complete, runnable **remote code-writing task** end to end: the environment runs on its own machine, training runs on the GPU machine, and the only thing connecting them is a URL. + +Code writing makes a good first agentic RL task for concrete reasons: the reward can be computed objectively from unit tests, so no judge model is needed; multi-turn interaction is genuinely useful (write → run → fix → submit); and OpenEnv's bundled code environment means almost no environment-side development. + +The runnable code lives in [`cookbook/rl/openenv_code/`](https://github.com/modelscope/twinkle/tree/main/cookbook/rl/openenv_code) (remote OpenEnv service) and [`cookbook/rl/agentenv/`](https://github.com/modelscope/twinkle/tree/main/cookbook/rl/agentenv) (Firecracker microVM). The two examples share the **same task, the same tool names and the same reward formula** — only the execution backend differs — so they can be read side by side. + +## 1. Pick an Execution Backend + +| | **OpenEnv embedded** | **OpenEnv server mode** | **AgentENV** | +|---|---|---|---| +| Adapter | `OpenEnv` + `EnvPool` | `OpenEnvClient` | `AgentEnv` | +| Isolation | None (same process) | Process / container | microVM (KVM) | +| Executor | Depends on the env package | smolagents AST interpreter (`coding_env`) | Real CPython | +| Special hardware | None | None | Requires `/dev/kvm` | +| Runs untrusted code | ❌ | ⚠️ Only within a controlled whitelist | ✅ | +| `assert` / files / `pip` | Depends on the env package | ❌ | ✅ | +| Deployment cost | Zero | One `uvicorn` command | Deploy the AgentENV control plane | + +Recommendations: + +- **Pure-compute environments** (board games, text games, scoring logic that is safe to evaluate in-process) → embedded `OpenEnv`, sharded with `EnvPool`. +- **Code execution / scaling the env independently of training / conflicting dependencies** → `OpenEnvClient` (the path this guide follows). +- **Genuinely untrusted code** (the model may write files, install packages, spawn processes) → AgentENV; see [Agentic RL with Sandbox Environments](./Agentic-RL-Sandbox.md). + +> The key trade-off: OpenEnv's `coding_env` uses smolagents' `LocalPythonExecutor`, which is an **AST interpreter, not an OS-level sandbox**. It is good enough to enforce "only these modules may be imported", but do not use it for genuinely adversarial code. If you need that level of isolation, use AgentENV. + +## 2. Environment Side: Serve OpenEnv as a Service + +On the environment machine (no GPU, no KVM, no Docker required): + +```bash +pip install openenv +pip install -e /path/to/OpenEnv/envs/coding_env # brings in smolagents + +cd cookbook/rl/openenv_code +sh serve.sh # 4 workers x 64 sessions = 256 concurrent sessions +``` + +`serve.sh` starts [`server_app.py`](https://github.com/modelscope/twinkle/blob/main/cookbook/rl/openenv_code/server_app.py) from this folder rather than upstream's `coding_env.server.app`. It makes three **necessary** deviations from upstream defaults — copying upstream verbatim will bite you: + +```python +class ConcurrentCodeEnv(PythonCodeActEnv): + # 1. Upstream leaves this False, which makes create_app(max_concurrent_envs > 1) + # raise ConcurrencyConfigurationError and caps the server at ONE session. + SUPPORTS_CONCURRENT_SESSIONS = True + + def __init__(self): + super().__init__() + self._configure() + + def reset(self, **kwargs): + # The parent's reset() rebuilds the executor and transform with upstream + # defaults, so re-apply our config afterwards. + observation = super().reset() + self._configure() + return observation + + def _configure(self) -> None: + # 2. Upstream authorises only `import json`, so math / collections — which + # many MBPP solutions need — would fail. + self._executor = PyExecutor(additional_imports=list(ALLOWED_IMPORTS)) + # 3. Upstream's create_safe_coding_transform() overwrites observation.reward + # with code-style heuristics (-1.0 for open( / import os, +0.1 for short + # code). This task's reward comes from unit tests, so a style score on + # the same channel is pure noise. + self.transform = None + +app = create_app(ConcurrentCodeEnv, CodeAction, CodeObservation, + env_name='twinkle_code_env', max_concurrent_envs=MAX_CONCURRENT_ENVS) +``` + +Flipping `SUPPORTS_CONCURRENT_SESSIONS` is safe here: `create_app` receives the **class** as a factory, so every WebSocket connection gets a fresh instance, and `PythonCodeActEnv.__init__` builds a private executor and state that share nothing. + +**Do the capacity math**: concurrent sessions = `WORKERS x MAX_CONCURRENT_ENVS`, and it must be ≥ `BATCH_SIZE x NUM_GENERATIONS` (4 x 8 = 32 with the defaults, against a server capacity of 256). Connections beyond capacity are rejected outright, which shows up as a subset of trajectories in the batch whose observations are all `Error:`. + +> The HTTP `/step` and `/reset` endpoints will **not** work for this. Each of those requests builds a fresh env and `close()`s it on the way out, losing all state — they exist for debugging and stateless use. Multi-turn episodes must go over WebSocket, which `OpenEnvClient` handles for you. + +## 3. Training Side: Three Pieces of Wiring + +### 1. Tools: What the Model Sees + +Only two tools are exposed ([`tools.py`](https://github.com/modelscope/twinkle/blob/main/cookbook/rl/openenv_code/tools.py)): + +- `run_python(code)` — sent to the server and executed in the session. **The session namespace persists across turns**, so the model can define a function in one turn and call it in the next. +- `submit_solution(code)` — never sent to the server. Registered via `register_tool` and handled **client-side**, recording the final source on the env for the training loop to score. + +```python +def _submit_solution(env: OpenEnvClient, arguments: Dict[str, Any]) -> str: + code = (arguments.get('code') or '').strip() + if not code: + return "Error: 'code' argument is required." + env.submitted_code = code + return 'Solution submitted.' + +def register_tools(env: OpenEnvClient) -> OpenEnvClient: + env.submitted_code = None + return env.register_tool(TOOL_SCHEMA[1], _submit_solution) +``` + +This is a general pattern: **"an action in the environment" and "bookkeeping the trainer needs" are different things**. Keep the latter in a local handler instead of polluting the environment protocol. + +The system prompt has to state the backend's semantics precisely, or the model will code against the wrong mental model. The two examples are **opposites** on exactly this point: + +- OpenEnv session: `The interpreter keeps its state between calls`, plus the module whitelist and a note that there is no file or network access. +- AgentENV: `Each call runs in a FRESH process, so every snippet must be self-contained`. + +### 2. Environment: One Session per Trajectory + +```python +def make_env() -> OpenEnvClient: + env = OpenEnvClient( + env_name=OPENENV_ENV_NAME, # 'coding_env'; client + Action classes auto-discovered + base_url=OPENENV_BASE_URL, # a load-balancer address works too + tools=[TOOL_SCHEMA[0]], # only run_python goes to the server + message_timeout_s=120.0, # code execution can be slow + ) + return register_tools(env) + + +def prepare_trajectories(samples, pool): + envs = [make_env() for _ in samples] + # reset() blocks on the WebSocket handshake plus server-side env creation, + # so run them concurrently — otherwise a batch of 32 serializes. + list(pool.map(lambda env: env.reset(), envs)) + + trajectories, tool_managers = [], [] + for sample, env in zip(samples, envs): + tool_managers.append(ToolManager(EnvTool.from_env(env))) + trajectories.append({ + 'messages': [ + {'role': 'system', 'content': SYSTEM_PROMPT}, + {'role': 'user', 'content': sample['prompt']}, + ], + 'tools': TOOL_SCHEMA, + }) + return trajectories, tool_managers, envs +``` + +Three things to get right: + +1. **Do not wrap `OpenEnvClient` in `EnvPool` / `@remote_class`.** The session's lifetime is owned by the server, so sharding it again through Ray buys nothing and only adds an RPC hop. +2. **One `ToolManager` per trajectory.** A `ToolManager` holds a specific env instance; sharing one would route every trajectory's tool calls into the same session. +3. **Close `envs` in a `finally`.** Sessions occupy server capacity, and leaking them makes later steps fail for want of capacity: + +```python +try: + all_trajectories = rollout(expand_prompts, tool_manager=tool_managers) + total_rewards, pass_rates = extract_rewards(envs, expanded, env_pool) +finally: + close_envs(envs, env_pool) +``` + +### 3. Reward: Use Unit Tests, Not a Judge Model + +Every MBPP sample ships a handful of `assert` statements. After the rollout, replay the hidden tests **in the same session**: + +```python +def run_tests(env, test_list, setup_code=''): + total = len(test_list) + solution = getattr(env, 'submitted_code', None) + if not solution: + return 0, total + if not _ok(env.execute({'code': solution})): # define the solution in the namespace + return 0, total + passed = 0 + for test in test_list: + expr = test.strip() + if expr.startswith('assert '): + expr = expr[len('assert '):] + result = env.execute({'code': f'print({expr})'}) + if _ok(result) and _last_line(result.observation.stdout) == 'True': + passed += 1 + return passed, total +``` + +Two deliberate design choices here: + +- **`env.execute()` instead of `env.step()`.** `execute()` returns the server's **raw** `StepResult`, so typed fields such as `exit_code` are readable; `step()` returns text rendered for the model and accumulates episode reward. Scoring logic should not show up in the model's conversation. +- **Rewriting `assert X` into `print(X)` and running one test per step**, rather than submitting one block of asserts. The sandbox is smolagents' AST interpreter, whose support for `assert` / `try` varies across versions while plain expressions are stable; and one step per test isolates a test that raises, so the rest still run (the reward uses the pass rate). + - Contrast: the AgentENV version runs real CPython, so it simply generates an ordinary script that wraps each assertion in `try/except` and prints `TESTS_PASSED n total`. Same reward, script shape dictated by backend capability. + +Reward shaping: + +```python +rate = passed / total if total else 0.0 +if total and passed == total: + return 1.0, rate # all tests pass +if getattr(env, 'submitted_code', None): + return 0.1 + 0.4 * rate, rate # submitted, partially correct -> 0.1 .. 0.5 +return 0.0, rate # never submitted +``` + +The shape guarantees "all correct > partially correct > submitted but all wrong > never submitted", and the 1.0 for a full pass is clearly above the 0.5 ceiling for partial credit, so the model cannot profit from shipping a fake implementation that only satisfies the first test. The **0.1 floor for submitting at all** exists to provide gradient signal early on: otherwise every trajectory scores 0, every GRPO group advantage is 0, and nothing is learned. + +### 4. GRPO: Group-Relative Advantages + +```python +batch = [dataset[(sample_cursor + i) % len(dataset)] for i in range(BATCH_SIZE)] +expanded = [s for s in batch for _ in range(NUM_GENERATIONS)] # N contiguous copies per problem +... +advantages = advantage_fn(total_rewards, num_generations=NUM_GENERATIONS, scale='group').tolist() +``` + +Note that `expanded` uses `for s in batch for _ in range(N)`: the N rollouts of one problem are **contiguous** in the list, which is the layout `GRPOAdvantage` slices groups by. Writing `for _ in range(N) for s in batch` misaligns every group — training appears to run but learns nothing. + +## 4. Running It + +Environment machine: + +```bash +cd cookbook/rl/openenv_code +sh serve.sh +``` + +Training machine (`pip install openenv` is enough; `coding_env` is not needed because the client class comes from `openenv`): + +```bash +cd cookbook/rl/openenv_code +OPENENV_BASE_URL=http://:8000 sh openenv_code_grpo.sh +``` + +Both environment variables and CLI arguments can be overridden at invocation: + +```bash +OPENENV_BASE_URL=http://10.0.0.5:8000 \ +MAX_TURNS=8 ENV_CONCURRENCY=32 \ +sh openenv_code_grpo.sh --batch-size 8 --num-generations 16 --max-steps 500 +``` + +After changing `batch-size` / `num-generations`, **scale server capacity to match**: `8 x 16 = 128` concurrent sessions still fits under `serve.sh`'s default 256, but beyond that raise `WORKERS` or `MAX_CONCURRENT_ENVS`. + +Key log line: + +``` +[Step 0] avg_reward=0.145, solve_rate=0.031, test_pass_rate=0.208, avg_turns=4.2 +``` + +| Metric | Meaning and interpretation | +|--------|----------------------------| +| `solve_rate` | Fraction passing every hidden test. This is the metric that must go up | +| `test_pass_rate` | Mean per-test pass rate. Smoother than `solve_rate`; watch it first early on | +| `avg_reward` | Mean shaped reward. If it rises while `solve_rate` does not, the model is farming the 0.1 submission floor | +| `avg_turns` | Mean turns. Pinned at `MAX_TURNS` means the model often burns its budget without submitting | + +## 5. Network Hardening (read before any shared deployment) + +**Neither OpenEnv nor AgentENV has any authentication.** AgentENV's README states it verbatim — *"AgentENV currently does not support authorization"* — and OpenEnv's `serve.sh` defaults to `--host 0.0.0.0`. Anyone who can reach the port can execute code in your sandbox and exhaust your capacity. + +Pick one, easiest first: + +| Situation | What to do | +|-----------|------------| +| Training on the same host (simplest) | `HOST=127.0.0.1 sh serve.sh` — never touches the network | +| Across hosts (recommended) | Security-group inbound rule allowing **only the training host's private IP/32**, or its security-group ID. **Never 0.0.0.0/0** | +| Stronger control needed | `iptables`/`nftables` source-IP restriction, or front it with nginx/caddy doing mTLS / token checks | + +⚠️ AgentENV's `aenv auth` API key is **not authentication** — any non-empty string works for a local deployment, since the CLI merely requires the field to be non-empty. Do not mistake it for a security control. + +For **egress** (where sandboxed code may connect), AgentENV enforces a node-level policy in `config/default.toml`, and the defaults are sensible: + +```toml +[network.egress] +# A sandbox's own egress policy cannot override these +always_denied_cidrs = [ + "10.0.0.0/8", "100.64.0.0/10", "127.0.0.0/8", + "169.254.0.0/16", # ← the critical one: blocks cloud metadata (169.254.169.254) + "172.16.0.0/12", "192.168.0.0/16", +] +``` + +That `169.254.0.0/16` entry matters most: it blocks the cloud metadata service, preventing sandboxed code from stealing temporary IAM credentials and escalating to your whole cloud account. The task in this guide (standard-library algorithm problems) needs no egress at all, so the defaults apply as-is. + +## 6. Troubleshooting + +| Symptom | Cause and fix | +|---------|---------------| +| Some trajectories' observations are all `Error:` | Server session capacity exhausted. Check `WORKERS x MAX_CONCURRENT_ENVS ≥ BATCH_SIZE x NUM_GENERATIONS` | +| `ConcurrencyConfigurationError` | You are serving upstream `coding_env.server.app` (single session). Use this folder's `server_app.py` | +| Odd ±0.1 / -1.0 values in the reward | The env's reward transform is still active; check `self.transform = None` | +| Model reports `import math` is not allowed | The `ALLOWED_IMPORTS` whitelist was lost; make sure `_configure()` is re-applied after `reset()` | +| Timeouts / message-wait errors | Raise `message_timeout_s`. Note the executor enforces three caps of its own: a 30s wall-clock limit per execution (`MAX_EXECUTION_TIME_SECONDS`), 10M operations, and 1M while-iterations — so `message_timeout_s` only helps above 30s | +| Batch skipped for too few valid trajectories | Most trajectories were filtered for length. Lower `MAX_TURNS` / `max_tokens`, or tighten tool output | +| Reward stays at 0 forever | Score a known-correct solution through `run_tests` in isolation first to prove the scoring path works, then suspect the model | + +## 7. Porting to Your Own Task + +Reusing this skeleton usually means changing four things: + +1. **Dataset**: replace `load_mbpp()` with your loader, producing `{'prompt', ...fields needed for scoring}`. +2. **Tools**: `TOOL_SCHEMA` plus handlers. Server-backed capabilities use the default action path; bookkeeping uses `register_tool`. +3. **Reward**: `extract_rewards()`. Prefer objectively computable signals (unit tests, exact match, executable verification); reach for a judge model only when there is none. +4. **System prompt**: it must describe the backend accurately (whether state persists across turns, which modules and permissions exist, the turn budget). + +The training loop, GRPO configuration, weight syncing and metrics normally need no changes. + +## Related Documents + +- Component reference: [Environments](../Components/Agentic/Envs.md) (`Env` abstraction, `EnvTool`, both OpenEnv modes, `EnvPool`) +- Deployment selection and cross-network access: [Agentic RL Deployment Guide](./Agentic-RL-Deployment.md) (backend matrix, four deployment tiers, SSH port forwarding) +- Multi-turn tool calling: [Multi-Turn Tool Usage](../Components/Agentic/Multi-Turn-Tool-Usage.md) +- Strong microVM isolation: [Agentic RL with Sandbox Environments](./Agentic-RL-Sandbox.md) +- Runnable examples: `cookbook/rl/openenv_code/` (remote OpenEnv), `cookbook/rl/agentenv/` (AgentENV microVM) +- OpenEnv upstream repository: diff --git a/docs/source_en/Usage Guide/Agentic-RL-Deployment.md b/docs/source_en/Usage Guide/Agentic-RL-Deployment.md new file mode 100644 index 000000000..1385318b3 --- /dev/null +++ b/docs/source_en/Usage Guide/Agentic-RL-Deployment.md @@ -0,0 +1,176 @@ +# Agentic RL Deployment Guide + +Twinkle assumes nothing about which sandbox you use or what network your machines sit in. This guide is a **lookup table**: pick an execution backend, then pick a deployment tier. + +The conclusion that shapes your planning, up front: + +> **The two dimensions are orthogonal.** Changing the execution backend means changing code (a different adapter class); changing the network topology means changing one environment variable. +> So **get the flow working on a single machine first — moving to multi-host or across the internet later requires no code changes at all.** + +## 1. First: Pick an Execution Backend + +Each of the three routes has a runnable example to copy from: + +| Your environment needs | Example | Where the env runs | Isolation | Memory per env | +|---|---|---|---|---| +| Pure compute (board games, text games, safely evaluable scoring) | [`cookbook/rl/multi_turn/`](https://github.com/modelscope/twinkle/tree/main/cookbook/rl/multi_turn) | In the training process, sharded across Ray workers by `EnvPool` | None | KBs | +| Code execution; env scaled independently of training | [`cookbook/rl/openenv_code/`](https://github.com/modelscope/twinkle/tree/main/cookbook/rl/openenv_code) | A standalone OpenEnv service (HTTP/WebSocket) | Process / container | KBs | +| Real execution semantics (`unittest`+mock, files, `pip`, subprocesses) | [`cookbook/rl/agentenv/`](https://github.com/modelscope/twinkle/tree/main/cookbook/rl/agentenv) | AgentENV Firecracker microVMs | microVM (KVM) | **~1GB** | + +The short version: + +- **Tasks solvable with the standard library** → OpenEnv. Three orders of magnitude lighter; no reason to reach for microVMs. +- **Running `unittest` + `@patch`, or the model needs to write files / install packages** → AgentENV is required. OpenEnv's `coding_env` uses smolagents' AST interpreter, where **decorators are silently ignored** (`decorator_list` is referenced nowhere in its source). Mocks stop working without raising, and the reward comes out as a plausible-looking wrong number. +- Full capability comparison: [Environments](../Components/Agentic/Envs.md). + +## 2. Second: Pick a Deployment Tier + +Four tiers, from single-machine validation to across the public internet. **L2 and L3 share identical training code** — only the network differs. + +| Tier | Situation | Environment deployment | Network | +|---|---|---|---| +| **L0** | Get it working | Same process / same host as training | None | +| **L1** | Single node, multi-GPU; move env overhead off the GPU process | `EnvPool` + dedicated CPU workers | None (inside the Ray cluster) | +| **L2** | Env scaled independently / dependency conflicts / isolation needed | Separate host | Private network + security group | +| **L3** | Env host and training host on different networks | Separate host | SSH port forwarding | + +### L0 — Single machine + +Embedded (the environment is instantiated in the driver, zero network overhead): + +```bash +cd cookbook/rl/multi_turn && python multi_turn_grpo.py +``` + +Or an OpenEnv server bound to loopback, with training on the same host: + +```bash +cd cookbook/rl/openenv_code +HOST=127.0.0.1 sh serve.sh & +OPENENV_BASE_URL=http://127.0.0.1:8000 sh openenv_code_grpo.sh \ + --batch-size 2 --num-generations 4 --max-steps 2 +``` + +**Prove the full path at low concurrency before scaling up** — it saves both time and money. But **do not shrink it too far**: the trajectory count (`batch-size × num-generations`) must be ≥ `--model-gpus` (4 by default), otherwise too few survive the length filter and the whole batch is skipped — the log just repeats `skipping this batch`, which looks like a hang rather than an error. The `2×4=8` above leaves 2x headroom. + +### L1 — Environments on dedicated CPU workers ("OpenEnv + Ray") + +Environments are still instantiated in-process, but `EnvPool` shards N instances onto a separate CPU `DeviceGroup`, off the GPU process's memory and GIL: + +```bash +cd cookbook/rl/multi_turn +ENV_REMOTE=1 ENV_NUM_WORKERS=8 ENV_POOL_SIZE=64 python multi_turn_grpo.py +``` + +| Variable | Effect | +|---|---| +| `ENV_REMOTE=1` | Put envs on a dedicated CPU DeviceGroup; unset runs them in the driver (zero RPC overhead) | +| `ENV_NUM_WORKERS` | Number of CPU workers; each rank becomes one `EnvPool` worker | +| `ENV_POOL_SIZE` | Pool capacity; `0` means auto (the trajectory count) | + +⚠️ **This tier applies only to the embedded `OpenEnv`.** Do not put `OpenEnvClient` or `AgentEnv` into an `EnvPool` — their session/sandbox lifetime is owned by the server, so sharding again through Ray buys nothing and only adds an RPC hop. + +### L2 — Environment on its own host (private network, the default for most setups) + +This is **the recommended default for most setups**: a company's GPU and CPU machines are usually already in the same VPC / datacentre / Kubernetes cluster, and need no VPN at all. + +Bind the environment to the **private NIC** (not `0.0.0.0`): + +```bash +# OpenEnv +HOST=10.0.1.20 sh serve.sh + +# AgentENV: configure its listener to the private IP after installation +``` + +The training side changes exactly one environment variable: + +```bash +OPENENV_BASE_URL=http://10.0.1.20:8000 sh openenv_code_grpo.sh # OpenEnv +AENV_API_URL=http://10.0.1.20:8000 sh agentenv_grpo.sh # AgentENV +``` + +Three things to get right: + +| Item | How | +|---|---| +| Narrow the inbound rule | The security group allows **only the training host's IP/32 or its security-group ID**, port 8000 only, never `0.0.0.0/0` | +| Confirm the bind took effect | `ss -tlnp \| grep 8000` must show `10.0.1.20:8000`, not `0.0.0.0:8000` | +| Supervise the service | Use [`deploy/openenv-server.service`](https://github.com/modelscope/twinkle/blob/main/cookbook/rl/deploy/openenv-server.service) — a rollout batch dies with the server | + +Traffic inside one VPC never touches the public internet, and with a source-restricted security group **this tier needs no extra networking components at all**. + +### L3 — Across networks (training and environment on different networks) + +The training commands are identical to L2; only the address changes to the local end of an SSH port forward — see the next section. + +## 3. Cross-Network Access: SSH Port Forwarding + +Twinkle supports exactly two ways to reach an environment: + +| Situation | Connection | +|---|---| +| Same host / same private network (L0–L2) | **Direct HTTP** | +| Training and environment on different networks (L3) | **SSH port forwarding** | + +The framework neither ships nor recommends VPN / NAT-traversal components. That layer belongs to your network team's existing infrastructure and is orthogonal to the training code — Twinkle only ever sees an `http://host:port`. + +**Why SSH for cross-network**: neither OpenEnv nor AgentENV has **any authentication**, so a reachable port means reachable arbitrary code execution. SSH's role here is not "connecting" but **putting authentication in front of a service that has none** — reusing your existing SSH keys and audited logins, on a channel most companies already approve (a bastion host is exactly this pattern). + +SSH port forwarding operates at the TCP layer and is fully transparent to WebSocket — verified: handshake, session state across messages, and 8 concurrent sessions over a single connection all work. + +```bash +# Environment side: loopback only, unreachable from the network +HOST=127.0.0.1 sh serve.sh + +# Training side: set up the forward, then connect to the local port +ssh -N -o ServerAliveInterval=30 -o ExitOnForwardFailure=yes \ + -L 8000:127.0.0.1:8000 user@env-host & +OPENENV_BASE_URL=http://127.0.0.1:8000 sh openenv_code_grpo.sh +``` + +Two addresses are easy to confuse: `OPENENV_BASE_URL` points at the **local entry** `127.0.0.1:8000`, while the `127.0.0.1:8000` inside `-L` is resolved **from the environment host's point of view**. Same for AgentENV via `AENV_API_URL`. + +**Long runs must be supervised**: if the forward drops, the whole rollout batch is lost. Do not leave it in an interactive shell — use autossh: + +```bash +autossh -M 0 -N -o ServerAliveInterval=30 -o ServerAliveCountMax=3 \ + -L 8000:127.0.0.1:8000 user@env-host +``` + +On throughput: all traffic shares one TCP connection. Each turn carries only a few KB of code and output, and 8 concurrent sessions showed no strain in testing; at hundreds of concurrent sessions, open several forwards on different local ports to spread the load. + +> When both sides already sit in the same VPC / datacentre / Kubernetes cluster, L2's direct HTTP is enough and SSH is unnecessary — that covers the large majority of corporate setups. + +## 4. Capacity Planning + +**OpenEnv server**: concurrent session ceiling = `WORKERS × MAX_CONCURRENT_ENVS`, which must be ≥ `BATCH_SIZE × NUM_GENERATIONS`. `serve.sh` defaults to `4 × 64 = 256` while the default training config needs only `4 × 8 = 32`, leaving ample headroom. Connections beyond capacity are rejected outright, showing up as trajectories whose observations are all `Error:`. + +**AgentENV**: memory is the only hard constraint. + +``` +required memory = concurrent trajectories × per-sandbox memory + 8GB (AgentENV + OS) +concurrent trajectories = BATCH_SIZE × NUM_GENERATIONS +``` + +With the example's `--memory-mb 1024`: the default 32 concurrent needs ~40GB, while a `2×4=8` concurrent smoke test needs ~16GB. **Do not buy a large machine for the trial phase.** + +AgentENV also has hard prerequisites (`/dev/kvm`, kernel 6.8+, `modprobe ublk_drv`, `CAP_SYS_ADMIN`) that container instances typically fail — see [Agentic RL with Sandbox Environments](./Agentic-RL-Sandbox.md). + +## 5. Pre-Flight Checklist + +- [ ] The service is **not** bound to `0.0.0.0` (verify with `ss -tlnp | grep 8000`) +- [ ] L2: security-group inbound allows only the training host, not `0.0.0.0/0` +- [ ] L3: the environment binds `127.0.0.1` and is reached only through SSH port forwarding +- [ ] Capacity ≥ `BATCH_SIZE × NUM_GENERATIONS` (plus the memory check for AgentENV) +- [ ] The environment service is supervised (systemd) and the SSH forward too (autossh), neither left in an interactive shell +- [ ] The full path was proven with `--batch-size 2 --num-generations 4 --max-steps 2` +- [ ] Trajectory count (`batch-size × num-generations`) ≥ `--model-gpus`, or batches are skipped with only a warning +- [ ] AgentENV's `always_denied_cidrs` was not loosened to "make the sandbox reachable" + +## Related Documents + +- End-to-end training flow and reward design: [Agentic RL Best Practices](./Agentic-RL-Best-Practices.md) +- AgentENV single-machine deployment in detail: [Agentic RL with Sandbox Environments](./Agentic-RL-Sandbox.md) +- Adapter and API reference: [Environments](../Components/Agentic/Envs.md) +- Deployment templates: `cookbook/rl/deploy/` diff --git a/docs/source_en/Usage Guide/Agentic-RL-Sandbox.md b/docs/source_en/Usage Guide/Agentic-RL-Sandbox.md new file mode 100644 index 000000000..eee282ca3 --- /dev/null +++ b/docs/source_en/Usage Guide/Agentic-RL-Sandbox.md @@ -0,0 +1,281 @@ +# Agentic RL with Sandbox Environments (AgentENV) + +This guide shows how to use [AgentENV](https://github.com/kvcache-ai/AgentENV) to give Twinkle **real sandbox environments** for Agentic RL training. Use it when the model must execute code, install packages, run tests or modify files inside a real operating system (tool-integrated reasoning, SWE-style tasks). + +Runnable example: `cookbook/rl/agentenv/`. + +## When You Need AgentENV + +Twinkle offers two kinds of execution environments; pick based on **whether the code is trusted**: + +| | OpenEnv (`EnvPool`) | AgentENV (`AgentEnv`) | +|---|---|---| +| Executes in | The Ray worker process | A dedicated Firecracker microVM | +| Isolation | None (shares the training process) | Hardware virtualization | +| Best for | Environment logic you wrote (game rules, graders) | **Untrusted model-generated code** | +| Install packages / mutate the OS | No | Yes, and it is wiped on destroy | +| Startup cost | Function call | ~50ms (snapshot restore) + HTTP RTT | + +If the model only emits an answer that your Python function grades, OpenEnv is enough and you can stop reading. **Once the model's output is code that gets executed, use AgentENV.** + +## Prerequisite Check + +AgentENV relies on KVM hardware virtualization, which imposes hard requirements on the host. Run this on the target machine: + +```bash +uname -r # need >= 6.8 +ls -l /dev/kvm # must exist and be read/write +modinfo ublk_drv >/dev/null && echo ublk-ok # need the ublk kernel module +``` + +All three must pass. Typical situations: + +- **Bare metal**: usually satisfied out of the box. +- **Cloud VMs / GPU instances**: require nested virtualization, which most providers disable by default. +- **Containers / K8s Pods / managed notebooks**: depend on the **host** kernel and `/dev/kvm`, plus privileged mode and a `/dev` mount. If the host lacks KVM, nothing inside the container can work around it. + +If the check fails, deploy the AgentENV server on a machine that passes and reach it over HTTP from the training side (see "Split Deployment"). + +## Single-Machine Topology + +The simplest setup runs the training process and the AgentENV server on the same host: + +``` +┌────────── One machine (bare metal, GPUs + KVM) ──────────┐ +│ │ +│ Training process (Ray/torchrun) agentenv server │ +│ ├─ model (GPU 0-3) ──HTTP──> 127.0.0.1:8000 │ +│ ├─ sampler (GPU 4-7) └─ Firecracker │ +│ └─ driver: rollout + tool calls sandboxes │ +└──────────────────────────────────────────────────────────┘ +``` + +A single node does not need AgentENV's gateway/scheduler (those are for multi-node clusters); talk to the server's `:8000` directly. + +**Split deployment** (GPU host lacks KVM): install AgentENV on another machine and point `AENV_API_URL` at it — no code changes. + +## Step 1: Deploy the AgentENV Server + +**Option A — install script (recommended, Ubuntu 24.04)** + +```bash +curl -fsSL https://raw.githubusercontent.com/kvcache-ai/AgentENV/main/scripts/install.sh \ + | sudo AENV_HOME_PATH=/data/aenv bash +sudo systemctl start aenv +``` + +The script creates a dedicated non-root `aenv` account (granted only `CAP_NET_ADMIN`/`CAP_SYS_ADMIN` plus kvm group access), loads the ublk module, downloads runtime assets (Firecracker, kernel), and registers a systemd service. `AENV_HOME_PATH` is the data directory (default `/var/lib/aenv`); point it at a large disk since image layers and snapshots live there. + +**Option B — Docker** + +```bash +curl -fsSL https://raw.githubusercontent.com/kvcache-ai/AgentENV/main/scripts/docker-setup.sh | sudo bash +docker run -d --privileged -v /dev:/dev -p 8000:8000 ghcr.io/kvcache-ai/aenv-server:latest +``` + +Docker here is only a **deployment vehicle**; sandboxes are still Firecracker microVMs, so host KVM is still required. + +**Verify** + +```bash +curl http://127.0.0.1:8000/health # expect 204 +``` + +## Step 2: Build an Environment Template + +A template is the sandbox's factory image: dependencies are pre-installed and frozen into a snapshot, which is what makes ~50ms boots possible. **Bake dependencies into the template** — installing them per episode at training time costs tens of seconds per trajectory and is not acceptable. + +Install the CLI and authenticate (the install script already ships `aenv`; use `install-cli.sh` on a remote machine): + +```bash +aenv auth +# AENV server URL [http://localhost:8000]: http://127.0.0.1:8000 +# API key: dummy # any non-empty string works for a local deployment +``` + +Write a Dockerfile (see `cookbook/rl/agentenv/Dockerfile`): + +```dockerfile +FROM python:3.11-slim +# Configure a mirror if your network requires one +# ENV PIP_INDEX_URL=https://pypi.org/simple +RUN pip install --no-cache-dir numpy sympy +WORKDIR /workspace +``` + +Build it: + +```bash +aenv build cookbook/rl/agentenv/Dockerfile -t twinkle-math --cpu-count 1 --memory-mb 1024 +aenv template watch # wait until ready +aenv template list # confirm the template exists +``` + +`aenv build` understands `FROM / RUN / ENV / WORKDIR / USER` (`ENTRYPOINT` becomes the start command; `EXPOSE / VOLUME / LABEL` are ignored). Rebuild only when the Dockerfile changes. + +You can also use an image as-is: `aenv pull ubuntu:22.04 --name ubuntu`. + +## Step 3: Wire Up the Training Side and Smoke-Test + +```bash +pip install e2b # AgentENV exposes an E2B-compatible API; reuse the official SDK +``` + +Validate the path with five lines before launching training: + +```python +from twinkle_agentic.envs import AgentEnv + +env = AgentEnv(template='twinkle-math', api_url='http://127.0.0.1:8000') +print(env.reset().observation) # boots a sandbox +print(env.step('run_command', {'command': 'python -c "print(6*7)"'}).observation) +print(env.step('read_file', {'path': '/etc/os-release'}).observation[:80]) +env.close() +``` + +`AgentEnv` is a stateless HTTP client implementing the standard `Env` interface (`reset`/`step`/`tools`/`close`). It deliberately does **not** use `@remote_class`: sandbox placement, load balancing, pause/resume and failover are all handled server-side by AgentENV, and Ray plays no part in environment scheduling. + +Common parameters: + +| Parameter | Description | +|---|---| +| `template` | Template name, i.e. the value of `aenv build -t` | +| `api_url` | Server or gateway URL; the `E2B_API_URL` env var also works | +| `sandbox_timeout` | Sandbox idle timeout (s). On expiry AgentENV pauses (not kills) the sandbox and auto-resumes on next access. **Must exceed the longest single trajectory** | +| `command_timeout` | Per-command timeout (s) | +| `setup_commands` | Commands run once after every reset | +| `include_default_tools` | Expose the built-in `run_command`/`write_file`/`read_file`; default `True` | + +## Step 4: Run Training + +```bash +cd cookbook/rl/agentenv +AENV_API_URL=http://127.0.0.1:8000 AENV_TEMPLATE=twinkle-math sh agentenv_grpo.sh +``` + +The example is multi-turn GRPO on GSM8K: the model writes Python in the sandbox to compute the result, then calls `submit`; the reward comes from answer correctness. Overridable environment variables: + +| Variable | Default | Description | +|---|---|---| +| `AENV_API_URL` | `http://127.0.0.1:8000` | AgentENV endpoint | +| `AENV_TEMPLATE` | `twinkle-math` | Template name | +| `SANDBOX_TIMEOUT` | `600` | Sandbox idle timeout (s) | +| `ENV_CONCURRENCY` | `16` | Driver threads used to create/kill sandboxes concurrently | +| `MAX_TURNS` | `6` | Max tool-calling turns per episode | + +Training hyperparameters (`--model-gpus`/`--batch-size`/`--num-generations`, ...) are CLI flags, consistent with the other cookbook examples. + +## Customizing Your Task + +Three places to change; the AgentENV server needs no modification. + +### 1. Tools (`tools.py`) + +The AgentENV server does **not** define "tools" — it provides capability primitives (arbitrary command execution, filesystem access, port proxying). "Tools" are a client-side concept, registered on `AgentEnv`: + +```python +# Command-template tool (most common) +env.register_command_tool( + {'type': 'function', 'function': { + 'name': 'run_tests', + 'description': 'Run the task test suite.', + 'parameters': {'type': 'object', + 'properties': {'test_file': {'type': 'string'}}, + 'required': ['test_file']}, + }}, + 'cd /workspace && pytest {test_file} -x -q') # formatted with tool arguments + +# Arbitrary Python handler +def submit(env, arguments): + env.submitted_answer = str(arguments.get('answer', '')).strip() + return f'Answer submitted: {env.submitted_answer}' + +env.register_tool({'type': 'function', 'function': {'name': 'submit', ...}}, submit) +``` + +A handler has signature `handler(env, arguments) -> str` (the return value becomes the observation) and can use `env.run_command(...)` and `env.sandbox` (the raw E2B handle, for PTY, file watching and other capabilities). Registering an existing name overrides the built-in tool; `include_default_tools=False` disables the built-ins entirely. + +For complex tools, **bake the implementation into the template** and keep the handler a one-liner: + +```dockerfile +COPY tools/search.py /opt/tools/search.py # installed at template build time +``` +```python +env.register_command_tool({...}, 'python /opt/tools/search.py {query}') +``` + +This gives tool implementations versioning, snapshot distribution and zero runtime cost. + +### 2. Rewards + +Sandboxes do not produce rewards (`AgentEnv.evaluate` returns zeros). Score after the rollout in the training loop, as `extract_rewards` does in the example: read the state your tool handler stashed on `env` (e.g. `env.submitted_answer`) and compare it with the ground truth. + +To grade inside the sandbox (e.g. running a test suite), have the tool emit structured output and parse it on the driver side: + +```python +out = env.run_command({'command': 'pytest -q --json-report --json-report-file=/tmp/r.json; cat /tmp/r.json'}) +``` + +### 3. Termination + +`MultiTurnRollout` terminates when **the model stops emitting tool calls** (or hits `max_turns` / the length cap); it does not read `EnvTool.done`. So the system prompt must explicitly ask the model not to call more tools after submitting, otherwise every episode runs to `max_turns`. + +## Capacity Planning + +Sandboxes consume **CPU and memory** only, and Firecracker does **not** support GPU passthrough — you cannot run GPU workloads inside a sandbox. + +Single-machine estimate: + +``` +concurrent sandboxes = batch_size × num_generations +memory needed ≈ concurrent sandboxes × the template's --memory-mb +``` + +The example defaults to `batch_size=4 × num_generations=8 = 32` concurrent sandboxes at 1GB each, i.e. ~32GB of host memory headroom (on top of GPU memory). If memory is tight, adjust in this order: + +1. Lower the template's `--memory-mb` (512MB suffices for many tasks) +2. Reduce `batch_size` +3. Rely on AgentENV's auto-pause: idle sandboxes return memory to the host + +Also watch CPU contention: sandboxes compete with dataloader/tokenizer work, so keep `ENV_CONCURRENCY` at or below the number of spare cores. + +## Security + +For single-machine personal use (server bound to `127.0.0.1`) the risk is low. Once **other people** can reach it: + +1. **AgentENV's control plane has no authorization**, and upstream explicitly states it must not be exposed to the public network. Any non-empty API key is accepted and there is no tenant isolation — anyone holding a sandbox id can read, write or destroy that sandbox, and `GET /sandboxes` lists everyone's sandboxes. +2. **Build your own auth/tenant/quota layer** in front of it, keep AgentENV on a private network, and expose only your own API (ideally task-level, so sandboxes never appear in the external contract). +3. **Restrict sandbox egress by default** so model-generated code cannot scan your intranet (SSRF) or abuse bandwidth. Pass a network policy at creation time: + +```python +# base_policy: Default | Allow | Deny +{'base_policy': 'Deny', + 'egress': {'allowed_domains': ['pypi.org'], + 'denied_cidrs': ['10.0.0.0/8', '172.16.0.0/12', '192.168.0.0/16', '169.254.0.0/16']}} +``` + +`169.254.0.0/16` must be denied (cloud metadata service). Each sandbox already gets its own network namespace and iptables isolation. + +4. Always set `sandbox_timeout` so failed trajectories cannot leave zombie sandboxes holding resources. + +## Troubleshooting + +| Symptom | Cause and fix | +|---|---| +| Server fails with `/dev/kvm is not accessible` | The runtime account is not in the kvm group, or the host has no KVM. Run `sudo server --setup-host --runtime-user aenv --runtime-group aenv` and restart the service | +| `ublk_drv is not loaded` | `sudo modprobe ublk_drv`; kernels older than 6.8 must be upgraded | +| `ImportError: AgentEnv requires the E2B SDK` | `pip install e2b` | +| Sandbox creation reports a missing template | Check the name with `aenv template list`; confirm the build is ready (`aenv template watch`) | +| `pip install` fails inside the sandbox | Egress blocked by the network policy, or no mirror configured; prefer pre-installing in the template | +| Sandbox becomes unavailable mid-trajectory | `sandbox_timeout` is shorter than the trajectory duration, so it was paused. Increase it | +| Batches are slow to start | Increase `ENV_CONCURRENCY`; make sure dependencies are baked into the template rather than installed at runtime | +| Training skips batches for too few valid trajectories | Most trajectories were filtered for length. Reduce `MAX_TURNS`/`max_tokens`, or trim tool output (observations are already truncated at 32K chars) | + +## Related Documents + +- End-to-end walkthrough and backend selection: [Agentic RL Best Practices](./Agentic-RL-Best-Practices.md) (includes the remote OpenEnv code task) +- Deployment selection and cross-network access: [Agentic RL Deployment Guide](./Agentic-RL-Deployment.md) (backend matrix, four deployment tiers, SSH port forwarding) +- Component reference: `Components/Agentic/Envs.md` (the `Env` abstraction, `EnvTool`, both OpenEnv modes) +- Multi-turn tool usage: `Components/Agentic/Multi-Turn-Tool-Usage.md` +- Runnable examples: `cookbook/rl/agentenv/` (AgentENV), `cookbook/rl/openenv_code/` (remote OpenEnv counterpart) +- AgentENV upstream docs: diff --git a/docs/source_en/index.rst b/docs/source_en/index.rst index 9ea9d2e6c..154d00c68 100644 --- a/docs/source_en/index.rst +++ b/docs/source_en/index.rst @@ -14,6 +14,9 @@ Twinkle DOCUMENTATION Usage Guide/Server and Client/index.rst Usage Guide/NPU-Support.md Usage Guide/Train-as-a-Service.md + Usage Guide/Agentic-RL-Sandbox.md + Usage Guide/Agentic-RL-Best-Practices.md + Usage Guide/Agentic-RL-Deployment.md Usage Guide/Introduction-with-Qwen3.5.md Usage Guide/Embedding-Training.md diff --git a/docs/source_zh/index.rst b/docs/source_zh/index.rst index 2a8c88aeb..57cc87346 100644 --- a/docs/source_zh/index.rst +++ b/docs/source_zh/index.rst @@ -14,6 +14,7 @@ Twinkle DOCUMENTATION 使用指引/服务端和客户端/index.rst 使用指引/NPU的支持.md 使用指引/训练服务.md + 使用指引/Agentic RL部署与训练.md 使用指引/Qwen3.5最佳实践.md 使用指引/Embedding训练.md diff --git "a/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/Agentic RL\351\203\250\347\275\262\344\270\216\350\256\255\347\273\203.md" "b/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/Agentic RL\351\203\250\347\275\262\344\270\216\350\256\255\347\273\203.md" new file mode 100644 index 000000000..fea32845d --- /dev/null +++ "b/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/Agentic RL\351\203\250\347\275\262\344\270\216\350\256\255\347\273\203.md" @@ -0,0 +1,547 @@ +# Agentic RL 部署与训练 + +Agentic RL 要做两件事：把**执行环境**部署起来，再把**训练**接上去。这两件事是正交的——换执行后端改一个环境变量 `CODE_RL_BACKEND`，换网络拓扑改一个 URL，训练代码一行都不用动。所以先在单机把流程跑通，之后再升级到跨机、跨网络。 + +可运行示例：[`cookbook/rl/code_rl/`](https://github.com/modelscope/twinkle/tree/main/cookbook/rl/code_rl)。任务是 MBPP 上的多轮代码生成 GRPO，一份 `train.py` 支持两个后端，任务/工具/奖励公式完全相同，只有执行环境不同，可以直接做对照实验。 + +## 选执行后端 + +| | **OpenEnv 嵌入式** | **OpenEnv server 模式** | **AgentENV** | +|---|---|---|---| +| 适配器 | `OpenEnv` + `EnvPool` | `OpenEnvClient` | `AgentEnv` | +| 环境跑在哪 | 训练进程内 | 独立的 HTTP/WebSocket 服务 | 独立的 Firecracker microVM | +| 隔离级别 | 无 | 进程 / 容器 | microVM（KVM 硬件虚拟化） | +| 执行器 | 取决于环境包 | smolagents AST 解释器（`coding_env`） | 真实 CPython | +| 单环境内存 | KB 级 | KB 级 | **1GB 级** | +| 文件 / `pip` / 子进程 | 取决于环境包 | 不支持 | 支持，销毁即清零 | +| 特殊硬件 | 无 | 无 | 需要 `/dev/kvm` | +| 部署成本 | 零 | 一条 `uvicorn` 命令 | 需部署 AgentENV 控制面 | + +- **纯计算型环境**（棋类、文本游戏、能在训练进程内安全求值的判分逻辑）→ 嵌入式 `OpenEnv`。 +- **要执行代码，且环境需要独立于训练扩缩容** → `OpenEnv` server 模式。轻三个数量级，标准库能解决的任务没必要上 microVM。 +- **要跑 `unittest` + `@patch`，或模型需要写文件、装包、开子进程** → 必须 AgentENV。OpenEnv 的 `coding_env` 走 smolagents 的 `LocalPythonExecutor`，那是一个 **AST 解释器，不是操作系统级沙箱**：**装饰器会被静默忽略**（其源码里对 `decorator_list` 没有任何引用），mock 失效且不报错，reward 会算出一个看似正常的错误数值。它足以约束"只允许 import 白名单模块"，但不要用它跑对抗性代码。 + +后端能力的完整对比见[执行环境](../组件/Agentic/Envs.md)。 + +--- + +## 第一部分：部署 + +### 一、OpenEnv + +#### 嵌入式模式（零部署） + +环境直接在训练进程里实例化，没有网络开销： + +```bash +cd cookbook/rl/multi_turn && python multi_turn_grpo.py +``` + +环境实例多、CPU 开销大时，用 `EnvPool` 把它们分片到独立的 CPU `DeviceGroup`，不占 GPU 进程的内存和 GIL： + +```bash +ENV_REMOTE=1 ENV_NUM_WORKERS=8 ENV_POOL_SIZE=64 python multi_turn_grpo.py +``` + +| 变量 | 作用 | +|---|---| +| `ENV_REMOTE=1` | 环境放到专属 CPU DeviceGroup；不设则在 driver 本地跑（零 RPC 开销） | +| `ENV_NUM_WORKERS` | CPU worker 数，每个 rank 一个 `EnvPool` worker | +| `ENV_POOL_SIZE` | 池容量，`0` 表示自动取轨迹数 | + +`EnvPool` 只适用于嵌入式 `OpenEnv`。不要把 `OpenEnvClient` 或 `AgentEnv` 放进 `EnvPool`——它们的 session / sandbox 生命周期在服务端，Ray 再分片一次不会有任何收益，只多一跳 RPC。 + +#### server 模式 + +在环境机器上（不需要 GPU、KVM、Docker）： + +```bash +pip install openenv + +# coding_env 是 OpenEnv 仓库里的子包（包名 openenv-coding_env），不在 PyPI 上， +# 只能从源码装。serve.sh 起的 server_app.py 从它 import PythonCodeActEnv 和 +# PyExecutor；smolagents 也是它的依赖带进来的。 +git clone https://github.com/huggingface/OpenEnv.git +pip install -e OpenEnv/envs/coding_env + +cd cookbook/rl/code_rl +sh serve.sh # 4 workers x 64 sessions = 256 并发 session +HOST=127.0.0.1 sh serve.sh # 训练在同一台机器时，绑回环不出网 +``` + +训练机只需要 `pip install openenv`，客户端类由它提供，不需要装 `coding_env`。 + +#### server_app.py 对上游默认值的三处修改 + +`serve.sh` 起的是本目录下的 [`server_app.py`](https://github.com/modelscope/twinkle/blob/main/cookbook/rl/code_rl/server_app.py)，不是上游的 `coding_env.server.app`。直接照抄 upstream 会踩坑： + +```python +class ConcurrentCodeEnv(PythonCodeActEnv): + # 1. 上游默认 False，会让 create_app(max_concurrent_envs > 1) 直接抛 + # ConcurrencyConfigurationError，服务端被限制成单 session。 + SUPPORTS_CONCURRENT_SESSIONS = True + + def reset(self, **kwargs): + # 父类 reset() 会用上游默认值重建 executor 和 transform，之后必须重新配一遍。 + observation = super().reset() + self._configure() + return observation + + def _configure(self) -> None: + # 2. 上游只授权 import json，MBPP 常用的 math / collections 会直接失败。 + self._executor = PyExecutor(additional_imports=list(ALLOWED_IMPORTS)) + # 3. 上游的 create_safe_coding_transform() 会用代码风格启发式覆盖 + # observation.reward（含 open( / import os 罚 -1.0，短代码奖 +0.1）。 + # 本任务的奖励来自单元测试，同一通道上的风格分只会是噪声。 + self.transform = None +``` + +打开 `SUPPORTS_CONCURRENT_SESSIONS` 是安全的：`create_app` 拿到的是**类**（作为 factory），每个 WebSocket 连接都会新建一个实例，executor 和 state 都是实例私有的。 + +只用 HTTP 的 `/step`、`/reset` 端点是不行的——OpenEnv 的这些端点每次请求都新建一个 env、返回后立刻 `close()`，状态全丢。多轮 episode 必须走 WebSocket，`OpenEnvClient` 已经处理好了。 + +#### server 模式的容量 + +``` +并发 session 上限 = WORKERS x MAX_CONCURRENT_ENVS # serve.sh 默认 4 x 64 = 256 +必须 ≥ BATCH_SIZE x NUM_GENERATIONS # 训练默认 4 x 8 = 32 +``` + +超容量的连接会被服务端直接拒绝，表现为一批 rollout 里部分轨迹的观测全是 `Error:`。改大 `--batch-size` / `--num-generations` 后记得同步扩 `WORKERS` 或 `MAX_CONCURRENT_ENVS`。 + +### 二、AgentENV + +[AgentENV](https://github.com/kvcache-ai/AgentENV) 用 Firecracker microVM 提供真实操作系统语义，适合 tool-integrated reasoning、SWE 类任务。 + +#### 前提条件自检 + +在目标机器上执行，三项都通过才能继续： + +```bash +uname -r # 需要 >= 6.8 +ls -l /dev/kvm # 必须存在且可读写 +modinfo ublk_drv >/dev/null && echo ublk-ok # 需要 ublk 内核模块 +``` + +- **裸金属机器**：通常直接满足。 +- **云上 VM / GPU 实例**：需要云厂商开启嵌套虚拟化，多数默认关闭。 +- **容器 / K8s Pod / DSW 等托管环境**：取决于**宿主机**内核与 `/dev/kvm`，且需要 privileged + 挂载 `/dev`。宿主机不满足时，容器内无法通过任何软件手段绕过。 + +不满足时：找一台满足条件的机器单独部署 AgentENV，训练侧把 `AENV_API_URL` 指过去即可，代码零改动。Firecracker **不支持 GPU 直通**，沙箱里跑不了 GPU 任务，所以环境机不需要 GPU。 + +#### Step 1：部署 server + +方式 A — 安装脚本（推荐，Ubuntu 24.04）： + +```bash +curl -fsSL https://raw.githubusercontent.com/kvcache-ai/AgentENV/main/scripts/install.sh \ + | sudo AENV_HOME_PATH=/data/aenv bash +sudo systemctl start aenv +``` + +脚本会创建专用的非 root `aenv` 账户（仅授予 `CAP_NET_ADMIN`/`CAP_SYS_ADMIN` 和 kvm 组权限）、加载 ublk 模块、下载 Firecracker/内核等运行时资源，并注册 systemd 服务。`AENV_HOME_PATH` 是数据目录（默认 `/var/lib/aenv`），镜像层和快照都放这里，建议指向大容量磁盘。它同时会装上 `aenv` CLI。 + +方式 B — Docker： + +```bash +curl -fsSL https://raw.githubusercontent.com/kvcache-ai/AgentENV/main/scripts/docker-setup.sh | sudo bash +docker run -d --privileged -v /dev:/dev -p 8000:8000 ghcr.io/kvcache-ai/aenv-server:latest +``` + +这里的 Docker 只是**部署载体**，沙箱依然是 Firecracker microVM，仍然需要宿主机的 KVM。 + +验证：`curl http://127.0.0.1:8000/health` 期望返回 204。 + +单机场景不需要 AgentENV 的 gateway / scheduler（那是多节点用的），直接连 server 的 `:8000`。 + +#### Step 2：构建环境模板 + +模板是沙箱的"出厂镜像"，把依赖预装并固化成快照，沙箱启动才能做到 ~50ms。**依赖必须烤进模板**，不要在训练时现场 `pip install`——每条轨迹重复装几十秒是不可接受的。 + +`aenv` 是 Rust 二进制（不是 pip 包）。方式 A 的安装脚本已经带上了；只想装 CLI 用： + +```bash +curl -fsSL https://raw.githubusercontent.com/kvcache-ai/AgentENV/main/scripts/install-cli.sh | sudo bash +``` + +`aenv` 纯客户端，认证到哪就在哪建模板，在训练机上跑 `aenv build` 也可以： + +```bash +aenv auth +# AENV server URL [http://localhost:8000]: http://127.0.0.1:8000 +# API key: dummy # 本地部署任意非空字符串即可 + +aenv build cookbook/rl/code_rl/Dockerfile -t twinkle-code --cpu-count 1 --memory-mb 1024 +aenv template watch # 等到 ready +aenv template list # 确认模板存在 +``` + +`aenv build` 支持 `FROM / RUN / ENV / WORKDIR / USER`（`ENTRYPOINT` 转为启动命令；`EXPOSE / VOLUME / LABEL` 被忽略）。只在 Dockerfile 变化时才需要重新构建。不加工直接用现成镜像：`aenv pull ubuntu:22.04 --name ubuntu`。 + +复杂工具的实现也建议烤进模板，工具 handler 只是一行调用： + +```dockerfile +COPY tools/search.py /opt/tools/search.py +``` + +这样工具实现有版本、随快照分发、运行时零开销。 + +#### Step 3：训练侧冒烟测试 + +```bash +pip install e2b # AgentENV 暴露 E2B 兼容 API，客户端复用官方 SDK +``` + +先用 5 行代码验证链路，不要直接跑训练： + +```python +from twinkle_agentic.envs import AgentEnv + +env = AgentEnv(template='twinkle-code', api_url='http://127.0.0.1:8000') +print(env.reset().observation) # 沙箱启动 +print(env.step('run_command', {'command': 'python -c "print(6*7)"'}).observation) +env.close() +``` + +`AgentEnv` 是无状态 HTTP 客户端，实现标准 `Env` 接口。它**不使用 `@remote_class`**：沙箱的放置、负载均衡、休眠唤醒、故障转移全部由 AgentENV 服务端负责，Ray 不参与环境调度。常用参数： + +| 参数 | 说明 | +|---|---| +| `template` | 模板名，即 `aenv build -t` 的值 | +| `api_url` | server 或 gateway 地址 | +| `sandbox_timeout` | 沙箱空闲超时（秒）。超时后自动 pause（不是 kill），下次访问自动唤醒。**必须大于单条轨迹最长耗时** | +| `command_timeout` | 单条命令超时（秒） | +| `setup_commands` | 每次 reset 后执行的初始化命令 | +| `include_default_tools` | 是否暴露内置的 `run_command`/`write_file`/`read_file`，默认 `True` | + +#### 内存预算 + +内存是唯一硬约束： + +``` +并发沙箱数 = BATCH_SIZE x NUM_GENERATIONS +需要内存 = 并发沙箱数 x 模板的 --memory-mb + 8GB（AgentENV 自身 + 系统） +``` + +按 `--memory-mb 1024`：默认 32 并发约需 40GB；先打通流程用 `2 x 4 = 8` 并发只需约 16GB。**试水阶段别买大机器。** 内存不足时的调整顺序：降低模板 `--memory-mb`（多数任务 512MB 够用）→ 减小 `batch_size` → 依赖自动休眠（空闲沙箱 pause 后内存归还宿主机）。 + +另外注意 CPU 争抢：沙箱与 dataloader/tokenizer 抢 CPU，`ENV_CONCURRENCY` 不宜超过空闲核数。 + +### 三、网络与安全 + +#### 两种连接方式 + +twinkle 只支持两种： + +| 场景 | 连接方式 | +|---|---| +| 同机 / 同一内网、同一 VPC | **HTTP 直连** | +| 训练机与环境机跨网络 | **SSH 端口转发** | + +框架不提供、也不建议在 twinkle 里引入 VPN / NAT 穿透类组件。那属于网络团队的既有基础设施，与训练代码无关——twinkle 侧看到的永远只是一个 `http://host:port`。 + +#### HTTP 直连（多数场景的主线） + +企业的 GPU 机与 CPU 机通常本来就在同一个 VPC / IDC / K8s 集群里，不需要任何额外网络组件。环境侧绑**内网网卡**（不是 `0.0.0.0`）： + +```bash +HOST=10.0.1.20 sh serve.sh # OpenEnv；AgentENV 同理把监听地址配成内网 IP +``` + +训练侧只改一个环境变量： + +```bash +OPENENV_BASE_URL=http://10.0.1.20:8000 sh run_openenv.sh # OpenEnv +AENV_API_URL=http://10.0.1.20:8000 sh run_agentenv.sh # AgentENV +``` + +| 事项 | 做法 | +|---|---| +| 收窄入方向 | 安全组只放行**训练机的 IP/32 或其安全组 ID**，端口只开 8000，绝不用 `0.0.0.0/0` | +| 确认绑定生效 | `ss -tlnp \| grep 8000`，看到的必须是 `10.0.1.20:8000` 而不是 `0.0.0.0:8000` | +| 服务保活 | 用 [`deploy/openenv-server.service`](https://github.com/modelscope/twinkle/blob/main/cookbook/rl/deploy/openenv-server.service)——rollout 批次会随服务一起死 | + +#### SSH 端口转发（跨网络） + +**OpenEnv 和 AgentENV 都没有任何认证**，端口可达就等于任意代码执行可达。SSH 在这里的作用不是"连接"，而是**替零认证的服务加一道认证**——复用已有的 SSH 密钥和登录审计，且它通常是企业里早已批准的运维通道。 + +SSH 端口转发工作在 TCP 层，对 WebSocket 完全透明——已实测：握手、跨消息 session 状态、8 个并发 session 复用一条连接，全部正常。 + +```bash +# 环境侧：只绑回环，网络上彻底不可达 +HOST=127.0.0.1 sh serve.sh + +# 训练侧：先建立端口转发，再连本地端口 +ssh -N -o ServerAliveInterval=30 -o ExitOnForwardFailure=yes \ + -L 8000:127.0.0.1:8000 user@env-host & +OPENENV_BASE_URL=http://127.0.0.1:8000 sh run_openenv.sh +``` + +两个地址容易搞混：`OPENENV_BASE_URL` 填的是**本地入口** `127.0.0.1:8000`；`-L` 里那个 `127.0.0.1:8000` 是**在环境机视角**解析的目标。AgentENV 同理，改 `AENV_API_URL`。 + +长期任务必须保活，转发一断整批 rollout 全废。别挂在交互式 shell 里： + +```bash +autossh -M 0 -N -o ServerAliveInterval=30 -o ServerAliveCountMax=3 \ + -L 8000:127.0.0.1:8000 user@env-host +``` + +所有流量走一条 TCP 连接。每轮交互只有几 KB 代码和输出，实测 8 并发无压力；上百并发时可开多条转发映射到不同本地端口分流。 + +#### 沙箱出口（egress） + +限制沙箱内代码能访问哪里，防止模型生成的代码扫内网（SSRF）或滥用带宽。AgentENV 在 `config/default.toml` 里有节点级强制策略，默认值就不错： + +```toml +[network.egress] +# 沙箱自身的 egress policy 无法覆盖这些 +always_denied_cidrs = [ + "10.0.0.0/8", "100.64.0.0/10", "127.0.0.0/8", + "169.254.0.0/16", # ← 关键：挡住云元数据服务 169.254.169.254 + "172.16.0.0/12", "192.168.0.0/16", +] +``` + +`169.254.0.0/16` 这条最重要：封掉云厂商元数据服务，防止沙箱内代码偷 IAM 临时凭证进而拿到整个云账号权限。标准库算法题根本不需要出网，默认值直接用。需要放开时按沙箱粒度传策略： + +```python +# base_policy: Default | Allow | Deny +{'base_policy': 'Deny', + 'egress': {'allowed_domains': ['pypi.org', 'mirrors.aliyun.com'], + 'denied_cidrs': ['10.0.0.0/8', '172.16.0.0/12', '192.168.0.0/16', '169.254.0.0/16']}} +``` + +另外：AgentENV 的 `aenv auth` 那个 API key **不是认证**，本地部署填任意非空字符串即可通过，也没有租户隔离——拿到任意 sandbox-id 就能读写、销毁该沙箱，`GET /sandboxes` 会列出所有人的沙箱。要对外提供服务，必须自建认证/租户/配额层，把 AgentENV 锁在私网，只暴露自己的上层 API。 + +--- + +## 第二部分：训练 + +以 `cookbook/rl/code_rl/` 为例。选「写代码」作为任务的原因很直接：奖励可以用单元测试客观计算，不需要裁判模型；多轮交互天然有意义（写 → 试跑 → 修 → 提交）。 + +### 四、训练侧接线 + +#### 1. 工具：模型看到什么 + +只暴露两个工具（`backends/openenv.py` / `backends/agentenv.py`）： + +- `run_python(code)` —— 在环境里执行代码。 +- `submit_solution(code)` —— **不发给服务端**，用 `register_tool` 在客户端本地处理，只把最终源码记在 env 上，供训练循环打分。 + +```python +def _submit_solution(env, arguments: Dict[str, Any]) -> str: + code = (arguments.get('code') or '').strip() + if not code: + return "Error: 'code' argument is required." + env.submitted_code = code + return 'Solution submitted.' +``` + +这是一个通用模式：**「模型的动作」和「训练需要的记账」是两件事**，后者放本地 handler，不要污染环境协议。 + +System prompt 必须准确描述后端语义，否则模型会按错误的心智模型写代码。两个后端在这一点上是**相反的**： + +- OpenEnv session：`The interpreter keeps its state between calls`——这一轮定义函数、下一轮可以直接调用；并列出可用模块白名单、声明无文件/网络访问。 +- AgentENV：`Each call runs in a FRESH process, so every snippet must be self-contained`。 + +AgentENV 侧还用 `include_default_tools=False` 关掉了内置的 `run_command`/`read_file`/`write_file`，让动作空间与任务精确对齐，奖励归因才干净。 + +工具集是客户端概念，AgentENV 服务端本身**不定义"工具"**，它只提供能力原语（任意命令执行、文件读写、端口代理）。除 `register_tool(schema, handler)` 外还有命令模板型的 `register_command_tool`： + +```python +env.register_command_tool( + {'type': 'function', 'function': { + 'name': 'run_tests', + 'description': '运行任务测试集。', + 'parameters': {'type': 'object', + 'properties': {'test_file': {'type': 'string'}}, + 'required': ['test_file']}}}, + 'cd /workspace && pytest {test_file} -x -q') # 用工具参数格式化 +``` + +handler 签名是 `handler(env, arguments) -> str`（返回值即 observation），内部可用 `env.run_command(...)` 和 `env.sandbox`（原始 E2B 句柄，可访问 PTY、文件 watch 等）。同名注册会覆盖内置工具。 + +#### 2. 环境：一条轨迹一个 env + +```python +def prepare_trajectories(samples, pool): + envs = [backend.make_env() for _ in samples] + # reset() 阻塞在网络上（WebSocket 握手，或启动一个沙箱），必须并发， + # 否则一个 batch 32 条会串行等待。 + list(pool.map(lambda env: env.reset(), envs)) + + trajectories, tool_managers = [], [] + for sample, env in zip(samples, envs): + tool_managers.append(ToolManager(EnvTool.from_env(env))) + trajectories.append({ + 'messages': [ + {'role': 'system', 'content': backend.SYSTEM_PROMPT}, + {'role': 'user', 'content': sample['prompt']}, + ], + 'tools': backend.TOOL_SCHEMA, + }) + return trajectories, tool_managers, envs +``` + +三个要点： + +1. **不要用 `EnvPool` / `@remote_class` 包 `OpenEnvClient` 或 `AgentEnv`**。session / sandbox 的生命周期在服务端，用 Ray 再分片一次没有收益，只多一跳 RPC。 +2. **每条轨迹一个 `ToolManager`**。`ToolManager` 持有具体的 env 实例，共享会让所有轨迹的工具调用打到同一个 session。 +3. **`envs` 必须在 `finally` 里关**。它们占着服务端容量，泄漏会让后续 step 因容量不足而失败： + +```python +try: + all_trajectories = rollout(expand_prompts, tool_manager=tool_managers) + total_rewards, pass_rates = extract_rewards(envs, expanded, env_pool) +finally: + close_envs(envs, env_pool) +``` + +`MultiTurnRollout` 的终止条件是**模型不再发出工具调用**（或达到 `MAX_TURNS` / 长度上限），它不读 `EnvTool.done`。所以要在 system prompt 里明确要求"提交后不要再调用工具"，否则回合会跑满 `MAX_TURNS`。 + +#### 3. 奖励：用单元测试，不要用裁判模型 + +MBPP 每条样本自带若干 `assert` 断言。rollout 结束后回放隐藏测试，两个后端的**脚本形态由后端能力决定，奖励公式完全相同**： + +- **OpenEnv**：在**同一个 session** 里逐条执行，并把 `assert X` 改写成 `print(X)`。这样能把"断言求值为 False"和"代码崩了"区分开（失败的 `assert` 抛出的异常与解法内部崩溃无法区分），且一条测试抛异常不影响其余测试继续跑。该执行器是支持 `assert`/`try` 的，这里是可诊断性的选择，不是能力绕行。 +- **AgentENV**：真实 CPython，直接生成一个普通脚本，每条断言用自己的 `try` 包住，最后打印 `TESTS_PASSED n total`。 + +OpenEnv 侧用 `env.execute()` 而不是 `env.step()`：`execute()` 返回服务端**原始** `StepResult`，可以读 `exit_code` 这类结构化字段；`step()` 返回的是渲染给模型看的文本，并且会计入 episode。打分逻辑不应该出现在模型的对话里。 + +奖励整形： + +```python +rate = passed / total if total else 0.0 +if total and passed == total: + return 1.0, rate # 全部通过 +if getattr(env, 'submitted_code', None): + return 0.1 + 0.4 * rate, rate # 提交了，部分通过 → 0.1 ~ 0.5 +return 0.0, rate # 没提交 +``` + +形状上保证「全对 > 部分对 > 提交但全错 > 没提交」，且全对的 1.0 明显高于部分对的上限 0.5，避免模型学会"交一个能过第一个测试的假实现"。**提交动作本身给 0.1 底分**，是为了在训练早期给出梯度信号——否则一开始所有轨迹奖励全 0，GRPO 的组内优势全是 0，学不到任何东西。 + +沙箱本身不产生 reward（`AgentEnv.evaluate` 默认返回 0），打分统一在训练循环里做。需要在沙箱内判分时，让工具输出结构化结果再在 driver 侧解析。 + +#### 4. GRPO：组内相对优势 + +```python +batch = [dataset[(sample_cursor + i) % len(dataset)] for i in range(BATCH_SIZE)] +expanded = [s for s in batch for _ in range(NUM_GENERATIONS)] # 同题连续 N 份 +... +advantages = advantage_fn(total_rewards, num_generations=NUM_GENERATIONS, scale='group').tolist() +``` + +`expanded` 用的是 `for s in batch for _ in range(N)`，同一道题的 N 条 rollout 在列表里**连续排列**，`GRPOAdvantage` 按这个布局切分组。写成 `for _ in range(N) for s in batch` 会让分组完全错位，训练看起来能跑但学不到东西。 + +### 五、跑起来 + +先用小并发打通全链路再放大，省时间也省钱： + +```bash +cd cookbook/rl/code_rl + +# OpenEnv +OPENENV_BASE_URL=http://127.0.0.1:8000 sh run_openenv.sh \ + --batch-size 2 --num-generations 4 --max-steps 2 + +# AgentENV +AENV_API_URL=http://127.0.0.1:8000 AENV_TEMPLATE=twinkle-code sh run_agentenv.sh \ + --batch-size 2 --num-generations 4 --max-steps 2 +``` + +但别压得太狠：轨迹数（`batch-size x num-generations`）必须 ≥ `--model-gpus`（默认 4），否则长度过滤后剩余不足，整批会被直接跳过——日志只反复刷 `skipping this batch`，看起来像卡住而不是报错。上面的 `2 x 4 = 8` 留了 2 倍余量。 + +放大到正式训练： + +```bash +OPENENV_BASE_URL=http://10.0.0.5:8000 MAX_TURNS=8 ENV_CONCURRENCY=32 \ +sh run_openenv.sh --batch-size 8 --num-generations 16 --max-steps 500 +``` + +可覆盖的环境变量： + +| 变量 | 默认 | 说明 | +|---|---|---| +| `MAX_TURNS` | `6` | 单回合最大工具调用轮数 | +| `ENV_CONCURRENCY` | `16` | driver 并发创建/销毁/打分的线程数 | +| `OPENENV_BASE_URL` | `http://127.0.0.1:8000` | OpenEnv 服务地址（也可以是负载均衡器地址） | +| `OPENENV_ENV_NAME` | `coding_env` | 环境包名，决定客户端 + Action 类 | +| `OPENENV_MESSAGE_TIMEOUT_S` | `120` | 单条消息超时 | +| `AENV_API_URL` | `http://127.0.0.1:8000` | AgentENV 地址 | +| `AENV_TEMPLATE` | `twinkle-code` | 模板名 | +| `SANDBOX_TIMEOUT` | `600` | 沙箱空闲超时（秒） | +| `AENV_COMMAND_TIMEOUT` | `60` | 沙箱内单条命令超时（秒） | + +训练超参（`--model-gpus`/`--batch-size`/`--num-generations` 等）走 CLI 参数，默认值在 `common_args.sh`——两个后端共用同一份，这样奖励的变化才能归因到执行后端而不是超参差异。 + +#### 指标 + +日志每步打一行，同时通过 `swanlab.log` 上报（`project='twinkle'`，实验名 `code-rl-`）： + +``` +[Step 0] {'train/code_acc': 0.031, 'train/test_pass_rate': 0.208, 'train/avg_reward': 0.145, ...} +``` + +| 指标 | 含义与解读 | +|---|---| +| `train/code_acc` | pass@1：全部隐藏测试通过的比例。这是真正要涨的指标 | +| `train/test_pass_rate` | 平均单测通过率。比 `code_acc` 平滑，早期先看它动没动 | +| `train/avg_reward` | 整形后的奖励均值。若它涨而 `code_acc` 不涨，说明模型在薅 0.1 的提交底分 | +| `train/avg_turns` | 平均轮数。贴着 `MAX_TURNS` 说明模型经常不提交就用完轮数 | +| `train/max_turns` / `train/min_turns` | 批内轮数分布。`max_turns` 顶到 `MAX_TURNS` 说明有轨迹被截断；`min_turns` 接近 1 说明常有轨迹第一轮就提交或放弃 | +| `train/approx_kl` | 新旧策略的 KL（Schulman K3 估计）。突然飙升是策略崩塌的前兆 | +| `train/clip_ratio` | PPO 裁剪触发比例（另有 `_low` / `_high` 分向）。长期偏高说明单步更新太激进 | +| `train/token_kl_max` / `train/token_ratio_max` | 单 token 的 KL / 概率比极值，用于定位崩塌 | +| `train/policy_confidence` | `exp(mean_logp)`，策略平均置信度 | + +这些 GRPO 指标由 `model.add_metric('GRPOMetric', is_training=True, epsilon=...)` 注册，`epsilon` 要与 `set_loss('GRPOLoss', epsilon=...)` 保持一致，否则 `clip_ratio` 统计的裁剪阈值与实际生效的不是同一个。`train/entropy` 需要 `GRPOLoss(entropy_coef > 0)` 才会出现——注意那会往 loss 里加 entropy bonus，**改变训练目标**，不要只为了看指标而打开。 + +### 六、迁移到你自己的任务 + +复用示例骨架时通常只需要改四处，训练循环、GRPO 配置、权重同步、指标都不用动： + +1. **数据集**：`load_mbpp()` 换成你的加载函数，产出 `{'prompt', ...打分所需字段}`。 +2. **工具**：`TOOL_SCHEMA` + handler。服务端能力走默认 action 通路，客户端记账用 `register_tool`。 +3. **奖励**：`extract_rewards()` 里的 `score()`。优先选能客观计算的信号（单测、精确匹配、可执行校验），实在没有再上裁判模型。 +4. **System prompt**：必须准确描述后端语义（状态是否跨轮保留、有哪些模块/权限、轮数预算）。 + +### 七、排查 + +| 现象 | 原因与处理 | +|---|---| +| 一批轨迹里部分观测全是 `Error:` | 服务端容量不足。核对 `WORKERS x MAX_CONCURRENT_ENVS ≥ BATCH_SIZE x NUM_GENERATIONS` | +| `ConcurrencyConfigurationError` | 起的是上游 `coding_env.server.app`（单 session）。用本目录的 `server_app.py` | +| 奖励里混进 ±0.1 / -1.0 的怪值 | 环境的 reward transform 没关掉，检查 `self.transform = None` | +| 模型报 `import math` 不被允许 | `ALLOWED_IMPORTS` 白名单未生效；确认 `reset()` 之后重新调了 `_configure()` | +| 超时 / 消息等待报错 | 调大 `OPENENV_MESSAGE_TIMEOUT_S`。执行器自身另有三重上限：单次执行 30s wall-clock、1000 万次操作、100 万次 while 迭代，所以该值要留在 30s 以上才有意义 | +| AgentENV server 启动报 `/dev/kvm is not accessible` | 运行账户不在 kvm 组，或宿主机无 KVM。执行 `sudo server --setup-host --runtime-user aenv --runtime-group aenv` 后重启服务 | +| 报 `ublk_drv is not loaded` | `sudo modprobe ublk_drv`；内核 < 6.8 需升级内核 | +| `ImportError: AgentEnv requires the E2B SDK` | `pip install e2b` | +| 创建沙箱报模板不存在 | `aenv template list` 确认名字；构建是否 ready（`aenv template watch`） | +| 沙箱内 `pip install` 失败 | 出口网络被策略拦截，或未配 pip 镜像源；建议改为在模板里预装 | +| 轨迹中途报沙箱不可用 | `SANDBOX_TIMEOUT` 小于轨迹耗时，被自动 pause。调大该值 | +| batch 启动很慢 | 调大 `ENV_CONCURRENCY`；确认依赖已烤进模板而非运行时安装 | +| 有效轨迹数不足被跳过 | 多数轨迹超长被过滤。调小 `MAX_TURNS`/`--max-tokens`，或收敛工具输出（观测已默认截断到 32K 字符）；也要确认轨迹数 ≥ `--model-gpus` | +| 奖励长期为 0 | 先单独跑 `run_tests` 对一个已知正确解打分，确认打分链路本身是通的，再怀疑模型 | +| `Address already in use` | 端口被占。`ss -tlnp \| grep 8000` 查占用者，或换 `PORT=8001 sh serve.sh` | + +### 上线检查清单 + +- [ ] 服务**没有**绑在 `0.0.0.0`（`ss -tlnp | grep 8000` 确认） +- [ ] 内网直连：安全组入方向只放行训练机，不是 `0.0.0.0/0` +- [ ] 跨网络：环境侧绑 `127.0.0.1`，只经 SSH 端口转发访问 +- [ ] 容量 ≥ `BATCH_SIZE x NUM_GENERATIONS`（AgentENV 再核内存） +- [ ] 环境服务有保活（systemd），SSH 转发有保活（autossh），不是挂在交互式 shell 上 +- [ ] 已用 `--batch-size 2 --num-generations 4 --max-steps 2` 打通过全链路 +- [ ] 轨迹数（`batch-size x num-generations`）≥ `--model-gpus` +- [ ] 没有为了"让沙箱联通"而放宽 AgentENV 的 `always_denied_cidrs` + +## 相关文档 + +- 组件参考：[执行环境](../组件/Agentic/Envs.md)（`Env` 抽象、`EnvTool`、OpenEnv 两种模式、`EnvPool`） +- 多轮工具调用：[多轮工具调用](../组件/Agentic/Multi-Turn-Tool-Usage.md) +- 部署模板：`cookbook/rl/deploy/` +- 可运行示例：`cookbook/rl/code_rl/`（代码任务，两个后端）、`cookbook/rl/multi_turn/`（嵌入式 OpenEnv） +- OpenEnv 上游仓库： +- AgentENV 官方文档： diff --git "a/docs/source_zh/\347\273\204\344\273\266/Agentic/Envs.md" "b/docs/source_zh/\347\273\204\344\273\266/Agentic/Envs.md" index 675f05baf..fd9e9cc11 100644 --- "a/docs/source_zh/\347\273\204\344\273\266/Agentic/Envs.md" +++ "b/docs/source_zh/\347\273\204\344\273\266/Agentic/Envs.md" @@ -97,41 +97,100 @@ env_tool = EnvTool( ) ``` -## OpenEnv +## OpenEnv：两种接入模式 -`OpenEnv` 将基于 WebSocket 的 [OpenEnv](https://github.com/OpenEnv) 环境服务器适配为同步的 Twinkle `Env`。 +[OpenEnv](https://github.com/meta-pytorch/OpenEnv) 的环境包同时提供「Environment 实现」和「EnvClient 客户端」，因此 Twinkle 提供两个适配器，对应两种截然不同的部署形态： + +| | **嵌入式** `OpenEnv` | **服务端** `OpenEnvClient` | +|---|---|---| +| 环境运行位置 | 训练进程内（直接实例化 Environment） | 独立的 OpenEnv 服务进程/容器 | +| 通信 | 无（本地函数调用） | WebSocket 长连接（一个连接 = 一个 session） | +| 隔离性 | 无，与训练进程共享内存和 GPU 节点 | 进程级/容器级，可部署在完全独立的机器 | +| 依赖 | 环境包需装在训练节点 | 环境包只需装在环境节点 | +| 扩展方式 | `EnvPool` 按 Ray worker 分片 | 服务端自身的并发 session + 多副本 | +| 适用场景 | 纯计算型轻量环境（棋类、文本游戏） | 代码执行、需要隔离或需独立扩缩容的环境 | + +> **不要**把 `OpenEnvClient` 放进 `EnvPool`：session 的生命周期在服务端，用 Ray 再分片一次不会带来任何收益，只会多一层 RPC。 + +### 模式一：嵌入式 `OpenEnv` + +绕过 OpenEnv 的 FastAPI 服务，在训练进程内直接构造 Environment，零网络开销。 ```python from twinkle_agentic.envs.openenv import OpenEnv env = OpenEnv( - base_url='http://localhost:8000', - env_cls='coding_env.CodingEnv', # 可选的类型化客户端 - env_kwargs={'message_timeout_s': 30}, - tool_schema=[...], # 可选的工具定义 + env_name='openspiel_env', # 环境包名，自动发现 Environment / Action 类 + env_kwargs={'game_name': 'blackjack'}, # 传给 Environment 构造函数 ) +result = env.reset() +result = env.step('play', {'action': 'hit'}) ``` -### 参数 +| 参数 | 类型 | 说明 | +|------|------|------| +| `env_name` | `str` | OpenEnv 环境包名（如 `'coding_env'`）。自动从 `.server` 中发现 `*Environment` 类，并从包的 `__all__` 中发现 `*Action` 类。 | +| `env_cls` | `str` 或 class | 显式指定 Environment 类（`'module:ClassName'`），与 `env_name` 二选一。 | +| `env_kwargs` | `Dict` | 传给 Environment 构造函数的参数。 | +| `action_cls` | `str` 或 class | 显式指定 Action 类；省略时从 `env_name` 自动发现。 | +| `action_mapper` | `Callable` | `(tool_name, arguments) -> action`。默认把工具参数直接作为 Action 的字段。 | + +### 模式二：服务端 `OpenEnvClient` + +先在环境节点上把 OpenEnv 环境跑成一个普通 HTTP/WebSocket 服务（不需要 Docker）： + +```bash +pip install openenv +pip install -e /path/to/OpenEnv/envs/coding_env +uvicorn coding_env.server.app:app --host 0.0.0.0 --port 8000 --workers 4 +``` + +然后在训练侧为每条轨迹创建一个客户端。每个实例都持有自己的 WebSocket session，服务端为它维护一个独立的 Environment 实例： + +```python +from twinkle_agentic.envs.openenv import OpenEnvClient + +env = OpenEnvClient( + env_name='coding_env', # 自动发现 EnvClient 子类 + Action 类 + base_url='http://10.0.0.5:8000', # 也可用 OPENENV_BASE_URL 环境变量 + message_timeout_s=120, # 环境要跑长耗时代码时调大 +) +env.reset() +result = env.step('run_python', {'code': 'print(1 + 1)'}) +print(result.observation) # '2' +env.close() +``` | 参数 | 类型 | 说明 | |------|------|------| -| `base_url` | `str` | 运行中的 OpenEnv 服务器 URL。 | -| `env_cls` | `str` 或 class | 类型化客户端的点分导入路径或类。`None` 使用 `GenericEnvClient`。 | -| `env_kwargs` | `Dict` | 传递给客户端构造函数的额外参数。 | -| `tool_schema` | `List[ToolInfo]` | 通过 `tools()` 暴露的工具定义。 | -| `action_mapper` | `Callable` | 自定义函数，将 `(tool_name, args)` 映射为发送给服务器的动作字典。 | +| `env_name` | `str` | OpenEnv 环境包名，从中自动发现 `EnvClient` 子类与 `*Action` 类。 | +| `env_cls` | `str` 或 class | 显式指定客户端类（`'module:ClassName'`），与 `env_name` 二选一。 | +| `base_url` | `str` | 服务地址，`http(s)://` 或 `ws(s)://` 均可（自动转换）。缺省读取 `OPENENV_BASE_URL`。也可以填负载均衡器地址。 | +| `action_cls` | `str` 或 class | Action 类；省略时自动发现。 | +| `action_mapper` | `Callable` | `(tool_name, arguments) -> action`，返回 Action 实例或字段字典。 | +| `tools` | `List[ToolInfo]` | 暴露给模型的工具 schema。默认是单个 `run_python(code)`，与 OpenEnv 代码类环境对齐。 | +| `reset_kwargs` | `Dict` | 转发给服务端 `reset()` 的参数（如 `repl_env` 的 `task_prompt` / `expected_answer`）。也可以按 episode 修改 `env.reset_kwargs` 属性。 | +| `connect_timeout_s` | `float` | WebSocket 连接超时，默认 10s。 | +| `message_timeout_s` | `float` | 单条消息超时，默认 120s。 | +| `client_kwargs` | `Dict` | 传给 OpenEnv 客户端构造函数的额外参数。 | + +补充能力： + +- `register_tool(tool_info, handler)`：注册一个**在客户端本地执行**的工具，不发往服务端。典型用途是 `submit_solution` 这类记账工具——把模型的答案存到 env 上，供训练循环后续打分。同名工具会覆盖默认工具。 +- `execute(action)`：直接发送动作并返回服务端**原始** `StepResult`，可以读取 `exit_code` 等结构化字段。训练循环用它在同一个 session 里追加执行单元测试。 +- `episode_reward` / `last_result` / `client`：累计奖励、上一次原始结果、底层同步客户端。 + +**容量与并发**：`OpenEnvClient` 内部调用 OpenEnv 客户端的 `.sync()`，每个实例拥有独立的后台事件循环，因此可以放在线程池里并发 reset/step。但服务端必须能容纳全部并发 session：环境类需声明 `SUPPORTS_CONCURRENT_SESSIONS = True`，且 `create_app(..., max_concurrent_envs=N)` 要够大，否则多出的连接会被拒绝。OpenEnv 自带的 `coding_env` 默认是**单 session**（`SUPPORTS_CONCURRENT_SESSIONS = False`），需要子类化后打开；`cookbook/rl/openenv_code/server_app.py` 给出了完整写法。 ### 与 Rollout 集成使用 +两种模式的下游用法完全相同： + ```python -from twinkle_agentic.envs.openenv import OpenEnv from twinkle_agentic.envs.env_tool import EnvTool from twinkle_agentic.tools.tool_manager import ToolManager from twinkle_agentic.rollout.api_multi_turn import APIMultiTurnRollout -# 设置环境 -env = OpenEnv(base_url='http://localhost:8000', tool_schema=[...]) env.reset() # 桥接到 ToolManager @@ -143,7 +202,36 @@ rollout = APIMultiTurnRollout(api=api, tool_manager=manager, max_turns=10) results = rollout(trajectories) ``` -### 实现自定义环境 +端到端的多轮 GRPO 训练示例见[Agentic RL 部署与训练](../../使用指引/Agentic%20RL部署与训练.md)。 + +## EnvPool：分布式环境池 + +`EnvPool` 是一个 `@remote_class`，把 `pool_size` 个**嵌入式** `OpenEnv` 实例按 Ray worker 分片。每个 worker 只管理 `pool_size // world_size` 个槽位，`reset`/`step` 通过 `remote_function` 自动路由到持有该槽位的 worker。 + +```python +from twinkle_agentic.envs.openenv import EnvPool + +pool = EnvPool( + pool_size=64, + device_mesh=mesh, + env_kwargs={'env_name': 'openspiel_env', 'env_kwargs': {'game_name': 'blackjack'}}, +) + +# 每个槽位包装成一个标准 Env，可直接交给 EnvTool / ToolManager +envs = pool.get_adapters(64) +env_tools = EnvTool.from_env(envs[0]) +``` + +| 方法 | 说明 | +|------|------| +| `reset(idx)` / `step(idx, tool_name, arguments)` | 操作单个槽位。 | +| `reset_batch(indices)` / `step_batch(indices, tool_names, arguments_list)` | 批量操作，一次 RPC 覆盖多个槽位，按 `indices` 顺序返回。 | +| `get_adapters(n)` | 在 driver 侧把前 `n` 个槽位包装为 `EnvPoolAdapter`（标准 `Env`）。 | +| `close()` | 关闭全部环境。 | + +`EnvPoolAdapter` 实现标准 `Env` 接口，把 `reset`/`step` 代理到对应 worker；`step` 出错时返回 `done=True` 并把错误写入 `info['error']`，避免单个环境异常拖垮整批 rollout。 + +## 实现自定义环境 ```python from twinkle_agentic.envs.base import Env, StepResult diff --git "a/docs/source_zh/\347\273\204\344\273\266/Agentic/Multi-Turn-Tool-Usage.md" "b/docs/source_zh/\347\273\204\344\273\266/Agentic/Multi-Turn-Tool-Usage.md" index 1c13bfbe0..8b94b2ed4 100644 --- "a/docs/source_zh/\347\273\204\344\273\266/Agentic/Multi-Turn-Tool-Usage.md" +++ "b/docs/source_zh/\347\273\204\344\273\266/Agentic/Multi-Turn-Tool-Usage.md" @@ -143,25 +143,26 @@ rollout = APIMultiTurnRollout(api=api, tool_manager=manager, max_turns=10) ## 使用 OpenEnv 环境 -连接远程 OpenEnv WebSocket 服务器： +连接一个运行中的 OpenEnv 服务（每个客户端实例对应服务端一个独立 session）： ```python -from twinkle_agentic.envs.openenv import OpenEnv +from twinkle_agentic.envs.openenv import OpenEnvClient from twinkle_agentic.envs.env_tool import EnvTool -env = OpenEnv( +env = OpenEnvClient( + env_name='coding_env', base_url='http://localhost:8000', - env_cls='coding_env.CodingEnv', - tool_schema=[{ + tools=[{ 'type': 'function', 'function': { - 'name': 'submit', - 'description': '提交代码解决方案。', + 'name': 'run_python', + 'description': '在远程解释器中执行 Python 代码。', 'parameters': { 'type': 'object', 'properties': { 'code': {'type': 'string'}, }, + 'required': ['code'], }, }, }], @@ -172,6 +173,8 @@ env_tools = EnvTool.from_env(env) manager = ToolManager(env_tools) ``` +如果希望环境跑在训练进程内（无隔离、零网络开销），把 `OpenEnvClient` 换成 `OpenEnv(env_name=..., env_kwargs=...)`。两种模式的完整对比见[执行环境](./Envs.md)。 + ## 每轨迹独立 ToolManager 当每个轨迹需要独立工具集时（例如，轨迹绑定的状态）： diff --git a/src/twinkle_agentic/envs/__init__.py b/src/twinkle_agentic/envs/__init__.py index a3cf38814..4633039c8 100644 --- a/src/twinkle_agentic/envs/__init__.py +++ b/src/twinkle_agentic/envs/__init__.py @@ -1,4 +1,5 @@ # Copyright (c) ModelScope Contributors. All rights reserved. +from .agentenv import AgentEnv from .base import Env, StepResult from .env_tool import EnvTool -from .openenv import EnvPool, EnvPoolAdapter, OpenEnv +from .openenv import EnvPool, EnvPoolAdapter, OpenEnv, OpenEnvClient diff --git a/src/twinkle_agentic/envs/agentenv.py b/src/twinkle_agentic/envs/agentenv.py new file mode 100644 index 000000000..47be62c46 --- /dev/null +++ b/src/twinkle_agentic/envs/agentenv.py @@ -0,0 +1,363 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""AgentENV adapter: thin client-side Env over an AgentENV (AENV) deployment. + +AgentENV (https://github.com/kvcache-ai/AgentENV) runs Firecracker microVM +sandboxes behind an E2B-compatible HTTP API. Unlike ``OpenEnv``/``EnvPool``, +this adapter deliberately does NOT use ``@remote_class``: sandbox placement, +load balancing, pause/resume and node failover are all handled server-side by +AgentENV's gateway/scheduler/orchestrator. The adapter is a stateless HTTP +client and can be instantiated directly inside rollout workers. + +Prerequisites (done once, outside training): + 1. Deploy the AgentENV server (single node) or gateway+scheduler cluster. + 2. Build a template, e.g. ``aenv pull ubuntu:22.04 --name my-env``. + 3. ``pip install e2b`` on the training side. + +Usage:: + + env = AgentEnv(template='my-env', api_url='http://gateway:8080') + result = env.reset() + result = env.step('run_command', {'command': 'echo hello'}) + env.close() +""" +import os +from typing import Any, Callable, Dict, List, Optional + +from twinkle.data_format import Trajectory +from twinkle.data_format.message import Tool as ToolInfo +from twinkle.utils import get_logger +from .base import Env, StepResult + +logger = get_logger() + +_MAX_OBSERVATION_CHARS = 32 * 1024 + +_DEFAULT_TOOLS: List[ToolInfo] = [ + { + 'type': 'function', + 'function': { + 'name': 'run_command', + 'description': 'Run a shell command inside the sandbox and return its output.', + 'parameters': { + 'type': 'object', + 'properties': { + 'command': { + 'type': 'string', + 'description': 'The shell command to execute.' + }, + 'cwd': { + 'type': 'string', + 'description': 'Working directory (optional).' + }, + }, + 'required': ['command'], + }, + }, + }, + { + 'type': 'function', + 'function': { + 'name': 'write_file', + 'description': 'Write text content to a file inside the sandbox.', + 'parameters': { + 'type': 'object', + 'properties': { + 'path': { + 'type': 'string', + 'description': 'Absolute file path in the sandbox.' + }, + 'content': { + 'type': 'string', + 'description': 'Text content to write.' + }, + }, + 'required': ['path', 'content'], + }, + }, + }, + { + 'type': 'function', + 'function': { + 'name': 'read_file', + 'description': 'Read a text file from the sandbox.', + 'parameters': { + 'type': 'object', + 'properties': { + 'path': { + 'type': 'string', + 'description': 'Absolute file path in the sandbox.' + }, + }, + 'required': ['path'], + }, + }, + }, +] + + +def _require_e2b(): + """Import the e2b SDK lazily with an actionable error message.""" + try: + from e2b import Sandbox + except ImportError as e: + raise ImportError('AgentEnv requires the E2B SDK to talk to an AgentENV server:\n' + ' pip install e2b\n' + 'Then point it at your deployment via api_url/api_key or the ' + 'E2B_API_URL / E2B_SANDBOX_URL / E2B_API_KEY environment variables.') from e + return Sandbox + + +def _truncate(text: str, limit: int = _MAX_OBSERVATION_CHARS) -> str: + if len(text) <= limit: + return text + return text[:limit] + f'\n... [truncated, {len(text) - limit} chars omitted]' + + +def _format_command_output(stdout: str, stderr: str, exit_code: int) -> str: + parts = [] + if stdout: + parts.append(stdout) + if stderr: + parts.append(f'[stderr]\n{stderr}') + if exit_code != 0: + parts.append(f'[exit code: {exit_code}]') + return _truncate('\n'.join(parts)) if parts else '(no output)' + + +class AgentEnv(Env): + """Env backed by one AgentENV sandbox per episode. + + Lifecycle mapping: + * ``reset`` -> kill the previous sandbox (if any) and create a fresh + one from ``template``; AgentENV's scheduler picks the node. + * ``step`` -> execute a tool inside the sandbox (sticky-routed to + the owning node via the sandbox id header, handled by the SDK). + * ``close`` -> kill the sandbox. + + Built-in tools (can be disabled via ``include_default_tools=False``): + ``run_command``, ``write_file``, ``read_file``. Task-specific tools can + be added with :meth:`register_tool` (arbitrary python handler) or + :meth:`register_command_tool` (shell command template), or by + subclassing. Tool errors never raise; they come back as observations so + the rollout loop can continue or let the model recover. + + Note: rewards are not produced by the sandbox. Keep the default + ``evaluate`` (zeros) and score trajectories with a separate reward + function, or subclass and override ``step``/``evaluate``. + """ + + def __init__(self, + template: str, + api_url: Optional[str] = None, + api_key: Optional[str] = None, + sandbox_timeout: int = 300, + command_timeout: int = 120, + setup_commands: Optional[List[str]] = None, + sandbox_envs: Optional[Dict[str, str]] = None, + metadata: Optional[Dict[str, str]] = None, + refresh_timeout: bool = True, + include_default_tools: bool = True, + **kwargs): + """ + Args: + template: AgentENV template name/ID (``aenv pull ... --name