diff --git a/.gitignore b/.gitignore index baa3010da..883fb9d58 100644 --- a/.gitignore +++ b/.gitignore @@ -56,6 +56,10 @@ venv/ ENV/ env.bak/ venv.bak/ +!unirl/rollout/env/ +!unirl/rollout/env/**/ +!unirl/rollout/env/**/*.py +!unirl/rollout/env/**/*.md # Jupyter / notebooks .ipynb_checkpoints/ diff --git a/README.md b/README.md index 08ff3e17b..f1e83f3d6 100644 --- a/README.md +++ b/README.md @@ -126,7 +126,7 @@ their `Sample` contains everything needed to resume, while stateful environment episodes and tool sessions must currently be dropped. Cross-worker stateful resume is deferred until its resource ownership and teardown contract is implemented. -See the [agent-loop guide](unirl/rollout/loop/README.md) for the environment, +See the [agent environment guide](unirl/rollout/env/README.md) for the environment, tool, trajectory, and partial-resume contracts. ## Getting Started ⚡ diff --git a/examples/alfworld/alfworld_grpo.yaml b/examples/alfworld/alfworld_grpo.yaml index b912444dd..83d9f98e2 100644 --- a/examples/alfworld/alfworld_grpo.yaml +++ b/examples/alfworld/alfworld_grpo.yaml @@ -108,7 +108,7 @@ rollout: cuda_graph_max_bs: 16 enable_lora: false env: - _target_: unirl.rollout.loop.alfworld_env.AlfworldEnv + _target_: unirl.rollout.env.alfworld.AlfworldEnv split: train max_steps: 15 # keep == config.max_turns step_penalty: 0.0 # binary terminal success; add a small penalty for efficiency later diff --git a/examples/alfworld/alfworld_grpo_async.yaml b/examples/alfworld/alfworld_grpo_async.yaml index 28c9d85cb..d49977c5f 100644 --- a/examples/alfworld/alfworld_grpo_async.yaml +++ b/examples/alfworld/alfworld_grpo_async.yaml @@ -114,7 +114,7 @@ rollout: cuda_graph_max_bs: 16 enable_lora: false env: - _target_: unirl.rollout.loop.alfworld_env.AlfworldEnv + _target_: unirl.rollout.env.alfworld.AlfworldEnv split: train max_steps: 15 # keep == config.max_turns step_penalty: 0.0 diff --git a/examples/alfworld/alfworld_grpo_partial.yaml b/examples/alfworld/alfworld_grpo_partial.yaml index 77cc9e049..0eb064fc6 100644 --- a/examples/alfworld/alfworld_grpo_partial.yaml +++ b/examples/alfworld/alfworld_grpo_partial.yaml @@ -107,7 +107,7 @@ rollout: cuda_graph_max_bs: 16 enable_lora: false env: - _target_: unirl.rollout.loop.alfworld_env.AlfworldEnv + _target_: unirl.rollout.env.alfworld.AlfworldEnv split: train max_steps: 15 # keep == config.max_turns step_penalty: 0.0 diff --git a/examples/deep_research/deep_research_search_judge.yaml b/examples/deep_research/deep_research_search_judge.yaml index dad9e22fd..2263eea05 100644 --- a/examples/deep_research/deep_research_search_judge.yaml +++ b/examples/deep_research/deep_research_search_judge.yaml @@ -119,12 +119,12 @@ rollout: cuda_graph_max_bs: 16 enable_lora: false env: - _target_: unirl.rollout.loop.tool_environment.ToolEnvironment + _target_: unirl.rollout.env.tool_environment.ToolEnvironment max_turns: 8 # keep == config.max_turns tools: - - _target_: unirl.rollout.loop.tools.search.SearchTool + - _target_: unirl.rollout.env.tools.search.SearchTool top_k: 10 - - _target_: unirl.rollout.loop.tools.visit.VisitTool + - _target_: unirl.rollout.env.tools.visit.VisitTool endpoint: ${oc.env:SUMMARY_URL,""} model: ${oc.env:SUMMARY_MODEL,""} episode_sampling: diff --git a/examples/deep_research/deep_research_search_judge_async.yaml b/examples/deep_research/deep_research_search_judge_async.yaml index f8371f00e..a0f513b8f 100644 --- a/examples/deep_research/deep_research_search_judge_async.yaml +++ b/examples/deep_research/deep_research_search_judge_async.yaml @@ -150,12 +150,12 @@ rollout: cuda_graph_max_bs: 16 enable_lora: false env: - _target_: unirl.rollout.loop.tool_environment.ToolEnvironment + _target_: unirl.rollout.env.tool_environment.ToolEnvironment max_turns: 8 # keep == config.max_turns tools: - - _target_: unirl.rollout.loop.tools.search.SearchTool + - _target_: unirl.rollout.env.tools.search.SearchTool top_k: 10 - - _target_: unirl.rollout.loop.tools.visit.VisitTool + - _target_: unirl.rollout.env.tools.visit.VisitTool endpoint: ${oc.env:SUMMARY_URL,""} model: ${oc.env:SUMMARY_MODEL,""} episode_sampling: diff --git a/examples/deep_research/deep_research_search_judge_partial.yaml b/examples/deep_research/deep_research_search_judge_partial.yaml index e9234aa17..80d0b5ada 100644 --- a/examples/deep_research/deep_research_search_judge_partial.yaml +++ b/examples/deep_research/deep_research_search_judge_partial.yaml @@ -137,12 +137,12 @@ rollout: cuda_graph_max_bs: 16 enable_lora: false env: - _target_: unirl.rollout.loop.tool_environment.ToolEnvironment + _target_: unirl.rollout.env.tool_environment.ToolEnvironment max_turns: 8 # keep == config.max_turns tools: - - _target_: unirl.rollout.loop.tools.search.SearchTool + - _target_: unirl.rollout.env.tools.search.SearchTool top_k: 10 - - _target_: unirl.rollout.loop.tools.visit.VisitTool + - _target_: unirl.rollout.env.tools.visit.VisitTool endpoint: ${oc.env:SUMMARY_URL,""} model: ${oc.env:SUMMARY_MODEL,""} episode_sampling: diff --git a/unirl/README.md b/unirl/README.md index a01f18c6c..ba512e14b 100644 --- a/unirl/README.md +++ b/unirl/README.md @@ -93,7 +93,7 @@ before applying the same reward, advantage, and train-stack contracts. - `types/README.md`: the `Sample` / `Part` contract and migration from the retired request/response API. - `config/README.md`: flat-recipe config — `require`/precision validators, `_target_` instantiation, cross-component contracts. - `rollout/README.md`: rollout modes, engines, and the `Sample` / `Part` generation flow. -- `rollout/loop/README.md`: agent-loop environments, tools, trajectories, and partial-resume behavior. +- `rollout/env/README.md`: agentic environments, tools, trajectories, and partial-resume behavior. - `train/readme.md`: train stack, FSDP backend, injection, EMA shadow. - `algorithms/README.md`: per-track loss algorithms. - `reward/README.md`: reward backends and custom scorers. diff --git a/unirl/distributed/group/handle.py b/unirl/distributed/group/handle.py index 31800a15d..bd3e7e5cc 100644 --- a/unirl/distributed/group/handle.py +++ b/unirl/distributed/group/handle.py @@ -321,9 +321,9 @@ class PendingHandleCall: """Future-like result of :meth:`Handle.launch_nowait`: launched, not yet collected. ``ready()`` probes without blocking; ``wait()`` blocks without collecting; - ``result()`` blocks if needed, then runs the rebind + collect half of - ``handle_fn`` and returns the method's collected value. The collect half - runs at most once — rebind registers GC finalizers on the result refs — so + ``result()`` blocks if needed, then runs the resolution phase of + ``handle_fn`` and returns the method's collected value. The resolution + phase runs at most once — rebind registers GC finalizers on the result refs — so a successful ``result()`` caches its value and later calls return it; a ``result()`` that raised may be retried. """ @@ -350,12 +350,8 @@ def result(self) -> Any: if self._consumed: return self._value handle = self._handle - results = ray.get(self._refs) - results = [ - handle._rebind_tree(r, handle.workers[i], worker_local=self._worker_local) for i, r in enumerate(results) - ] _, _, collect_fn, _ = handle._method_configs[self._method_name] - self._value = collect_fn(handle, results) + self._value = handle._resolve_call(collect_fn, self._refs, worker_local=self._worker_local) self._consumed = True return self._value @@ -605,8 +601,9 @@ def _make_handle_fn( grad_mode and call_id are passed as dedicated parameters to Worker.call (not via kwargs) so dispatch internals remain unaware of grad state. - The non-blocking twin is :meth:`launch_nowait` + - :meth:`PendingHandleCall.result` below — keep the halves in parity. + Both this blocking form and :meth:`launch_nowait` + + :meth:`PendingHandleCall.result` are thin sequencing over the shared + :meth:`_launch_call` / :meth:`_resolve_call` phases. """ def handle_fn(*args, **kwargs): @@ -627,36 +624,22 @@ def handle_fn(*args, **kwargs): call_id = f"{method_name}_{next(self._grad_call_counter)}" input_metas = collect_leaves(args, TensorRef) + collect_leaves(tuple(kwargs.values()), TensorRef) - batch_size = infer_batch_size(args, kwargs) - # Only DP_SCATTER/DP_SCATTER_HEAD split the per-sample batch by dp_size, so only - # they require divisibility; BROADCAST/SCATTER must not be rejected (main #202). - if ( - dispatch_mode in (Dispatch.DP_SCATTER, Dispatch.DP_SCATTER_HEAD) - and batch_size is not None - and batch_size % self.dp_size != 0 - ): - raise ValueError(f"batch_size={batch_size} not divisible by dp_size={self.dp_size}") - - shards = dispatch_fn(self, args, kwargs, batch_size) - # Locality + cross-worker transfer is the transport's policy: its - # localize makes every ref resolvable on its dst worker (GLOBAL = - # identity; worker-local = NCCL/IPC routing). It needs controller - # topology + per-shard dst identity, passed directly. - transport_cls = self.pool.transport_cls - worker_local = issubclass(transport_cls, WorkerLocalTransport) - shards = transport_cls.localize(shards, self.pool, self.device_ids, self.worker_ids) - # grad_mode/call_id passed as dedicated args, not mixed into kwargs - refs = execute_fn(method_name, shards, grad_mode=ctx is not None, call_id=call_id) - results = ray.get(refs, timeout=ray_get_timeout) - - # Rebind before collect: results[i] comes from workers[i], - # so worker attribution is unambiguous at this point. For worker-local - # this registers the decref GC finalizer; GLOBAL lifecycle is - # queue-managed, so skip rebind/GC there. - results = [self._rebind_tree(r, self.workers[i], worker_local=worker_local) for i, r in enumerate(results)] - - # Collect: merge primary rank results - collected = collect_fn(self, results) + refs, worker_local = self._launch_call( + method_name, + dispatch_mode, + dispatch_fn, + execute_fn, + args, + kwargs, + grad_mode=ctx is not None, + call_id=call_id, + ) + collected = self._resolve_call( + collect_fn, + refs, + worker_local=worker_local, + ray_get_timeout=ray_get_timeout, + ) if ctx is not None: output_metas = collect_leaves(collected, TensorRef) @@ -676,26 +659,28 @@ def handle_fn(*args, **kwargs): handle_fn.__doc__ = f"SPMD handle: {method_name} (dispatch={dispatch_fn.__name__})" return handle_fn - # ── Non-blocking launch ── - - def launch_nowait(self, method_name: str, *args, **kwargs) -> PendingHandleCall: - """Launch a @distributed method without blocking: the dispatch → localize → - execute half of ``handle_fn``, stopping before ``ray.get``. + # ── Shared call phases ── - Always ``grad_mode=False`` / ``call_id=None`` (a pending call is never - valid under a GradContext, so the ``_grad_call_counter`` single-thread - assumption is untouched). Kept in line-parity with ``handle_fn`` above — - same divisibility gate, same localize. ``result()`` on the returned - :class:`PendingHandleCall` runs the collect half. + def _launch_call( + self, + method_name: str, + dispatch_mode: Dispatch, + dispatch_fn: Callable, + execute_fn: Callable, + args: tuple, + kwargs: dict, + *, + grad_mode: bool, + call_id: Optional[str], + ) -> Tuple[List, bool]: + """Launch a distributed call; returns ``(refs, worker_local)``. + + Runs dispatch → localize → execute for both the blocking + ``handle_fn`` and the non-blocking :meth:`launch_nowait`. """ - try: - dispatch_mode, dispatch_fn, _, execute_fn = self._method_configs[method_name] - except KeyError: - raise AttributeError( - f"{method_name!r} is not a @distributed method of {_owning_class(self.role_cls).__name__}" - ) from None - batch_size = infer_batch_size(args, kwargs) + # Only DP_SCATTER/DP_SCATTER_HEAD split the per-sample batch by dp_size, so only + # they require divisibility; BROADCAST/SCATTER must not be rejected (main #202). if ( dispatch_mode in (Dispatch.DP_SCATTER, Dispatch.DP_SCATTER_HEAD) and batch_size is not None @@ -704,10 +689,67 @@ def launch_nowait(self, method_name: str, *args, **kwargs) -> PendingHandleCall: raise ValueError(f"batch_size={batch_size} not divisible by dp_size={self.dp_size}") shards = dispatch_fn(self, args, kwargs, batch_size) + # Locality + cross-worker transfer is the transport's policy: its + # localize makes every ref resolvable on its dst worker (GLOBAL = + # identity; worker-local = NCCL/IPC routing). It needs controller + # topology + per-shard dst identity, passed directly. transport_cls = self.pool.transport_cls worker_local = issubclass(transport_cls, WorkerLocalTransport) shards = transport_cls.localize(shards, self.pool, self.device_ids, self.worker_ids) - refs = execute_fn(method_name, shards, grad_mode=False, call_id=None) + # grad_mode/call_id passed as dedicated args, not mixed into kwargs + refs = execute_fn(method_name, shards, grad_mode=grad_mode, call_id=call_id) + return refs, worker_local + + def _resolve_call( + self, + collect_fn: Callable, + refs: List, + *, + worker_local: bool, + ray_get_timeout: Optional[float] = None, + ): + """Resolve a launched call into its collected method return value. + + Runs ray.get → rebind → collect for both the blocking + ``handle_fn`` and :meth:`PendingHandleCall.result`. + """ + results = ray.get(refs, timeout=ray_get_timeout) + # Rebind before collect: results[i] comes from workers[i], + # so worker attribution is unambiguous at this point. For worker-local + # this registers the decref GC finalizer; GLOBAL lifecycle is + # queue-managed, so skip rebind/GC there. + results = [self._rebind_tree(r, self.workers[i], worker_local=worker_local) for i, r in enumerate(results)] + # Collect: merge primary rank results + return collect_fn(self, results) + + # ── Non-blocking launch ── + + def launch_nowait(self, method_name: str, *args, **kwargs) -> PendingHandleCall: + """Launch a @distributed method without blocking: the launch phase of + ``handle_fn``, stopping before ``ray.get``. + + Always ``grad_mode=False`` / ``call_id=None`` (a pending call is never + valid under a GradContext, so the ``_grad_call_counter`` single-thread + assumption is untouched). ``result()`` on the returned + :class:`PendingHandleCall` runs the resolution phase. + """ + try: + dispatch_mode, dispatch_fn, _, execute_fn = self._method_configs[method_name] + except KeyError: + raise AttributeError( + f"{method_name!r} is not a @distributed method of {_owning_class(self.role_cls).__name__}" + ) from None + + refs, worker_local = self._launch_call( + method_name, + dispatch_mode, + dispatch_fn, + execute_fn, + args, + kwargs, + grad_mode=False, + call_id=None, + ) return PendingHandleCall(self, method_name, refs, worker_local) # ── Execute strategies ── diff --git a/unirl/rollout/engine/agentic/engine.py b/unirl/rollout/engine/agentic/engine.py index 66a299139..ccb8422f9 100644 --- a/unirl/rollout/engine/agentic/engine.py +++ b/unirl/rollout/engine/agentic/engine.py @@ -3,7 +3,8 @@ The engine is a ``BaseRolloutEngine`` on a DP-replicated slab. Each worker builds its **own local** inner single-turn engine + environment (the ``ComposedRolloutEngine`` build-inner pattern) and runs a **pool of drain threads** (``per_worker_concurrency`` of them) that pull single-trajectory tasks -from a central queue and run each as a plain synchronous multi-turn agent loop on its own thread: +from a central queue and run each through the configured harness (the multi-turn agent loop, +:class:`~unirl.rollout.harness.tool_agent.ToolAgentHarness`) on its own thread: per-turn ``inner.generate`` (the inner engine is safe for concurrent callers — its backend keeps the in-flight requests batching together) and ``env.step`` (the env is re-entrant across threads). @@ -59,6 +60,8 @@ from unirl.distributed.group.dispatch import Dispatch, Execute, distributed from unirl.rollout.engine.agentic.config import AgenticRolloutEngineConfig from unirl.rollout.engine.synchronous import BaseRolloutEngine, SyncRolloutEngine +from unirl.rollout.harness.protocol import HarnessContext, RolloutHarness +from unirl.rollout.harness.tool_agent import ToolAgentHarness from unirl.types.sample import Sample, _part_with_field from unirl.types.sampling import total_samples_per_prompt @@ -132,6 +135,16 @@ def __init__( # so siblings keep the GPU busy while one waits on a slow tool. self._concurrency = int(config.per_worker_concurrency) + # Task-internal control flow is the harness's (worker-side plugin); + # this runtime keeps queueing/concurrency/buffers/abort. Step 1 of the + # harness migration constructs the one existing harness here from the + # same config fields; config-selected harnesses come later. + self._harness: RolloutHarness = ToolAgentHarness(env=self._env, sampling=self._sp, max_turns=self._max_turns) + self._harness_ctx = HarnessContext( + engines={"policy": self._inner.generate}, + suspend=lambda: self._stopping, + ) + # Coordinator state (populated on rank 0 only, by set_workers). self._workers: List[Any] = [] self._role: str = "" @@ -359,63 +372,35 @@ def _drain_worker(self, coordinator: Any, role_name: str) -> None: raise def _run_one(self, task: Sample) -> Tuple[Sample, bool]: - """One trajectory's agent loop, on this drain thread. Returns ``(sample, done)``. - - ``done=True`` — terminal (the env said done, or ``max_turns`` reached). - ``done=False`` — **checkpointed** at a turn boundary because ``self._stopping`` - (partial rollout): the trajectory is carried and resumed next drive. - - Resume-aware: ``turns_done = len(gen_parts())`` (a carried partial continues - from where it stopped; ``env.reset`` is idempotent/turn-derived). The - ``_stopping`` check is at the **top of each turn**, so the in-flight turn - finishes naturally first (turn boundary, no mid-turn abort). Failure-isolated: - never raises into the drain. + """One trajectory, on this drain thread: delegate to the harness. Returns ``(sample, done)``. + + ``done=True`` — terminal (``completed``/``failed`` outcome). + ``done=False`` — **checkpointed** at a harness-chosen safe point + (``suspended``, partial rollout): carried and resumed next drive. + + The harness owns the task-internal control flow (turn loop, stop + conditions, teardown) and returns ``failed`` for task-level faults; + this runtime keeps the last-resort net (a harness BUG must not sink + the drain) and the tensor-side jobs: attaching an env-sourced reward, + and NaN-marking failures so an infrastructure fault never enters GRPO + as a legitimate low-scoring sibling (trainers give NaN zero advantage). + Failure-isolated: never raises into the drain. """ - sample = task - env_reward: Optional[float] = None try: - sample = self._env.reset(task) # [input(1)], root id = prompt id - turns_done = len(sample.gen_parts()) - for _ in range(self._max_turns - turns_done): - if self._stopping: # partial rollout: checkpoint at the turn boundary, carry & resume - return sample, False - # Concurrent-caller safe: fork() builds a fresh gen Part per call and - # the shared self._sp is only READ (resolve_sampling is pure) — an - # inner engine that mutated part.sampling_params in generate would - # turn this into a cross-thread race. - sample = self._inner.generate(sample.fork(1, sampling_params=self._sp)) # +[gen(1)] - observation, done, info = self._env.step(sample) # blocking tool boundary, own thread - # Env-sourced reward (LIN-519): interactive envs (ALFWorld, …) return a - # per-trajectory return in ``info["reward"]`` (last value = the episode - # return); tool-only envs (calculator/search) omit it — a no-op here. - if isinstance(info, dict) and info.get("reward") is not None: - env_reward = float(info["reward"]) - if done: - return self._attach_env_reward(sample, env_reward), True - if observation is not None: - sample = sample.observe(observation) # +[obs(1)] - return self._attach_env_reward(sample, env_reward), True # max_turns reached = terminal - except Exception as exc: # noqa: BLE001 — isolate: one bad trajectory must not sink the drain - # Mark the trajectory FAILED (NaN) instead of letting an infrastructure - # fault — backend outage, tool timeout, context overflow — enter GRPO as a - # legitimate low-scoring sibling. The trainers exclude NaN from the group's - # mean/std and give it zero advantage, so a failure neither rewards nor - # penalizes. Any partial ``env_reward`` is deliberately dropped: a reward - # collected before the fault does not describe a complete trajectory. - # Still ``done=True`` — the drain must not stall on a bad trajectory. - logger.warning("AgenticRolloutEngine: trajectory failed, marking failed: %s", exc, exc_info=True) - return self._attach_env_reward(sample, float("nan")), True - finally: - # Guaranteed teardown (LIN-533): end any open tool sessions / episodes for - # this trajectory — on success, crash, AND abort. Duck-typed like - # ``tool_schemas`` so envs without ``close`` are unaffected, and wrapped so - # a teardown error can never re-raise (``_run_one`` must not raise). - close = getattr(self._env, "close", None) - if close is not None: - try: - close(sample) - except Exception: # noqa: BLE001 — teardown must not sink the drain - logger.warning("AgenticRolloutEngine: env.close failed during teardown", exc_info=True) + outcome = self._harness.run(task, self._harness_ctx) + # Exhaustive over the outcome contract: a plugin returning an unknown + # status (or a malformed outcome object) must not slip into training + # as "completed" — it falls through to the except and is NaN-marked. + if outcome.status == "completed": + return self._attach_env_reward(outcome.sample, outcome.env_reward), True + if outcome.status == "suspended": + return outcome.sample, False + if outcome.status == "failed": + return self._attach_env_reward(outcome.sample, float("nan")), True + raise ValueError(f"unknown harness outcome status: {outcome.status!r}") + except Exception as exc: # noqa: BLE001 — harness bug; the partial trace is lost, the drain survives + logger.warning("AgenticRolloutEngine: harness outcome failed, marking failed: %s", exc, exc_info=True) + return self._attach_env_reward(task, float("nan")), True @staticmethod def _attach_env_reward(sample: Sample, reward: Optional[float]) -> Sample: diff --git a/unirl/rollout/engine/asynchronous.py b/unirl/rollout/engine/asynchronous.py index a930f9fd8..8dccd1c20 100644 --- a/unirl/rollout/engine/asynchronous.py +++ b/unirl/rollout/engine/asynchronous.py @@ -346,6 +346,7 @@ def __init__(self, rollout: Any, *, group_size: int, start_gen_id: int = 0) -> N self._buffer: VersionedBuffer[List["Sample"]] = VersionedBuffer() self._gen_id = int(start_gen_id) self._weight_version = 0 + self._drive_live = False @property def weight_version(self) -> int: @@ -356,7 +357,22 @@ def bump_weight_version(self) -> int: return self._weight_version def submit(self, tasks: List["Sample"]) -> None: - """Fire a background drive over a flat task list (fresh siblings + carried partials).""" + """Fire a background drive over a flat task list (fresh siblings + carried partials). + + Enforced double-pull guard: two live drains would double-pull the + coordinator queue, so a second ``submit`` before ``finalize_if_drained`` + reported the drive done (or before ``quiesce``) raises instead of + silently corrupting the drive. + """ + if self._drive_live: + raise RuntimeError( + "AsyncAgenticRolloutEngine.submit: prior drive still live — wait for " + "finalize_if_drained() to report it done or quiesce() first (a second " + "drain would double-pull the coordinator queue)." + ) + # Fail closed: the coordinator may start a drive before its call reports + # an error, so only finalize_if_drained() or quiesce() may re-arm submit. + self._drive_live = True self._rollout.submit(tasks) def poll(self) -> int: @@ -369,6 +385,7 @@ def finalize_if_drained(self) -> Optional[int]: completed = self._rollout.finalize_if_drained()[0] if completed is None: return None + self._drive_live = False return self._ingest(completed) def drain_freshest(self, n: int, *, max_staleness: int) -> Optional[List[List["Sample"]]]: @@ -384,6 +401,7 @@ def quiesce(self) -> List["Sample"]: version they completed under.""" carried = self._rollout.abort()[0] self.poll() + self._drive_live = False return carried def discard_roots(self, roots: Iterable[str]) -> int: @@ -406,8 +424,27 @@ def _ingest(self, completed: List["Sample"]) -> int: return len(completed) +def launch_ceiling(rollout_id: int, *, sync_interval: int, max_staleness: int, num_rollouts: int) -> int: + """The batch trainers' on-policy launch clamp — trainer POLICY, defined once. + + A generation launched now is consumed later, so how far ahead the gen_id + allocator may run is bounded to ``max_staleness`` weight-sync windows: + ``max_staleness=0`` ⇒ never launch into a future sync-window ⇒ no + generation crosses a sync ⇒ ``ratio≈1`` (on-policy). + + OWNERSHIP: this is trainer-side POLICY, not engine surface — its vocabulary + (``rollout_id`` / ``sync_interval`` / ``num_rollouts``) is the trainers', + the engine classes never call it, and it must never become an engine + method. It is hosted in this module only because it is the two batch + trainers' one shared torch-free home; the step loops that use it stay in + the trainers as visible statement order. + """ + return min(num_rollouts, ((rollout_id // sync_interval) + 1 + max_staleness) * sync_interval) + + __all__ = [ "AsyncAgenticRolloutEngine", "AsyncBatchRolloutEngine", + "launch_ceiling", "root_of", ] diff --git a/unirl/rollout/loop/README.md b/unirl/rollout/env/README.md similarity index 88% rename from unirl/rollout/loop/README.md rename to unirl/rollout/env/README.md index 960948223..f9d4adc45 100644 --- a/unirl/rollout/loop/README.md +++ b/unirl/rollout/env/README.md @@ -1,14 +1,13 @@ -# Agentic rollout loop +# Agentic environments and tools -This package defines the synchronous turn loop and the environment/tool contracts used by -agentic rollout. The distributed runtime is +This package defines the environment/tool contracts used by agentic rollout. The turn LOOP +itself lives in [`unirl/rollout/harness/tool_agent.py`](../harness/tool_agent.py) +(`ToolAgentHarness` — worker-side task control flow); the distributed runtime hosting it is [`AgenticRolloutEngine`](../engine/agentic/engine.py); trajectory storage is described by the [`Sample`/`Part` types](../../types/README.md). ## Contracts -- `RolloutEnginePort.generate(sample) -> Sample` fills one generation frontier. The production - agentic engine requires its inner engine to implement `SyncRolloutEngine`. - `Environment.reset(request) -> Sample` performs per-trajectory setup and may augment or replace the request. - `Environment.step(sample) -> (observation, done, info)` consumes the latest generated action. @@ -34,9 +33,10 @@ Sample -> fork -> blocking generate -> environment.step +------> observe -> next turn ``` -`AgentLoop.run` calls `environment.reset` once, uses -`total_samples_per_prompt(sampling_params)` for the first fork, and uses branch one for subsequent -turns. It stops when the environment returns `done=True` or after `max_turns`. +`ToolAgentHarness.run` calls `environment.reset` once, forks a one-sample continuation per +turn on the `"policy"` engine, and stops when the environment returns `done=True`, after +`max_turns`, or — at a turn boundary — when the runtime requests a cooperative suspension +(`HarnessContext.suspend_requested`). `AgenticRolloutEngine` is the production coordinator. Rank 0 expands `P` prompts into `P * n` single-trajectory tasks. Each worker runs up to `per_worker_concurrency` synchronous trajectories @@ -109,8 +109,8 @@ return through `info["reward"]`. GRPO siblings select the same game but run sepa - `max_turns` is the hard engine bound. When an environment also exposes `max_turns`, the production engine requires it to match. `ToolEnvironment` also terminates when no row calls a tool; `AlfworldEnv` terminates on simulator completion or `max_steps`. -- `AgentLoop` itself is blocking and does not call `Environment.close`; callers using it directly - own teardown. `AgenticRolloutEngine` calls `close` from `finally` on success, failure, and abort. +- `ToolAgentHarness.run` calls `Environment.close` from `finally` on success, failure, AND + suspension; a teardown error is logged, never raised. Current runnable configurations live under [`examples/deep_research`](../../../examples/deep_research) and [`examples/alfworld`](../../../examples/alfworld). diff --git a/unirl/rollout/env/__init__.py b/unirl/rollout/env/__init__.py new file mode 100644 index 000000000..98f6b86e7 --- /dev/null +++ b/unirl/rollout/env/__init__.py @@ -0,0 +1,21 @@ +"""Environments and tools for agentic rollout (LIN-492). See ``unirl/rollout/env/README.md``. + +:class:`~unirl.rollout.env.protocol.Environment` is the world side of a turn; +:class:`~unirl.rollout.env.tool_environment.ToolEnvironment` (with its +:mod:`~unirl.rollout.env.tools`) is the first concrete environment. The loop that +drives an environment lives in :mod:`unirl.rollout.harness.tool_agent` (worker-side, +hosted by :class:`~unirl.rollout.engine.agentic.engine.AgenticRolloutEngine`). +""" + +from unirl.rollout.env.protocol import Environment +from unirl.rollout.env.tool_environment import ToolEnvironment, parse_tool_call +from unirl.rollout.env.tools import CalculatorTool, StatefulTool, Tool + +__all__ = [ + "Environment", + "ToolEnvironment", + "parse_tool_call", + "Tool", + "StatefulTool", + "CalculatorTool", +] diff --git a/unirl/rollout/loop/alfworld_env.py b/unirl/rollout/env/alfworld.py similarity index 99% rename from unirl/rollout/loop/alfworld_env.py rename to unirl/rollout/env/alfworld.py index 4fbbaa9a8..72786b60d 100644 --- a/unirl/rollout/loop/alfworld_env.py +++ b/unirl/rollout/env/alfworld.py @@ -4,7 +4,7 @@ benchmark: a household task where the agent issues text commands (``go to shelf 1``, ``take mug 1``, ``clean mug 1 with sinkbasin 1`` …) and the simulator returns the next observation, ending with a binary task-success reward. Unlike the stateless -:class:`~unirl.rollout.loop.tool_environment.ToolEnvironment`, each trajectory is its +:class:`~unirl.rollout.env.tool_environment.ToolEnvironment`, each trajectory is its own **episode with evolving state** and the **reward comes from the simulator**. Fits the engine's ``reset(sample)->Sample`` / ``step(sample)->(obs, done, info)`` diff --git a/unirl/rollout/loop/environment.py b/unirl/rollout/env/protocol.py similarity index 69% rename from unirl/rollout/loop/environment.py rename to unirl/rollout/env/protocol.py index 8a5152bb1..8ec24a5be 100644 --- a/unirl/rollout/loop/environment.py +++ b/unirl/rollout/env/protocol.py @@ -1,7 +1,7 @@ -"""Environment — the world side of an agent-loop turn (LIN-492). +"""Environment — the world side of an agentic rollout turn (LIN-492). -See ``unirl/rollout/loop/README.md``. A structural ``Protocol`` seam only; concrete -environments (tool / critic) are designed separately. The loop treats it as optional. +See ``unirl/rollout/env/README.md``. A structural ``Protocol`` seam consumed by +rollout harnesses; concrete environments (tool / critic) are designed separately. """ from __future__ import annotations @@ -12,7 +12,7 @@ class Environment(Protocol): - """The world side of a turn. Optional for the loop; concrete environments are separate.""" + """The world side of a harness turn; concrete environments are separate.""" def reset(self, request: Sample) -> Sample: """Optional per-episode setup; return the (possibly augmented) request Sample.""" @@ -30,11 +30,11 @@ def step(self, sample: Sample) -> Tuple[Optional[Primitive], bool, dict]: ... def close(self, sample: Sample) -> None: - """Optional guaranteed teardown (LIN-533), called from the engine's ``finally`` on every - path — success, crash, and abort. The engine invokes it via ``getattr(env, "close", None)``, + """Optional guaranteed teardown (LIN-533), called from the harness's ``finally`` on every + path — success, crash, and abort. The harness invokes it via ``getattr(env, "close", None)``, so an env holding no per-trajectory resource need not implement it (default: no-op). Stateful envs use it to release handles exactly once — tool sessions - (:meth:`~unirl.rollout.loop.tool_environment.ToolEnvironment.close`) or ALFWorld episodes/ + (:meth:`~unirl.rollout.env.tool_environment.ToolEnvironment.close`) or ALFWorld episodes/ pooled templates. Must be idempotent and **must not raise**. """ ... diff --git a/unirl/rollout/loop/tool_environment.py b/unirl/rollout/env/tool_environment.py similarity index 94% rename from unirl/rollout/loop/tool_environment.py rename to unirl/rollout/env/tool_environment.py index b69942533..9518e8d83 100644 --- a/unirl/rollout/loop/tool_environment.py +++ b/unirl/rollout/env/tool_environment.py @@ -2,7 +2,7 @@ Turns the policy's ``{...}`` into a tool execution and feeds the result back as the next turn's observation, ending the episode when the model stops calling tools (it -gave a final answer) or ``max_turns`` is hit. See ``unirl/rollout/loop/README.md`` and the +gave a final answer) or ``max_turns`` is hit. See ``unirl/rollout/env/README.md`` and the real-world references it mirrors (relax ``DeepeyesEnv``, slime ``Geo3kEnv``): the loop stays mechanical; the *decision* — parse → execute → done — lives here. @@ -21,7 +21,7 @@ from typing import Any, Dict, List, Optional, Sequence, Tuple from uuid import uuid4 -from unirl.rollout.loop.tools.tool import StatefulTool, Tool +from unirl.rollout.env.tools.base import StatefulTool, Tool from unirl.types.primitives import Texts from unirl.types.sample import Primitive, Sample, _part_with_field @@ -106,7 +106,7 @@ def parse_tool_call(text: str) -> Optional[Dict[str, Any]]: class ToolEnvironment: """Agentic :class:`Environment`: parse tool calls, run tools, observe results, stop on a final answer. - Drives :class:`~unirl.rollout.loop.agent_loop.AgentLoop` over a frontier of one-or-more samples + Driven by :class:`~unirl.rollout.harness.tool_agent.ToolAgentHarness` over a frontier of one-or-more samples (the GRPO group / continuations). :meth:`step` parses each frontier sample's text, executes any tool call, and returns a row-aligned observation. The batch's ``done`` is True once **no** sample emits a tool call (all gave a final answer) or ``max_turns`` is reached. @@ -137,7 +137,7 @@ def reset(self, request: Sample) -> Sample: instance serves many concurrent trajectories on a worker). Stateful tools (LIN-533): mint a per-trajectory ``session_id`` (``uuid4``) for each - :class:`~unirl.rollout.loop.tools.tool.StatefulTool`, call ``session_start`` (cheap — the + :class:`~unirl.rollout.env.tools.base.StatefulTool`, call ``session_start`` (cheap — the handle opens lazily in :meth:`step`), and stamp the ids into the root Part's *control* bag under ``"tool_sessions"`` so :meth:`step`/:meth:`close` recover them position-independently across the fork/observe chain. ``uuid4`` avoids collisions between the ``n`` GRPO siblings @@ -221,9 +221,9 @@ def _run(self, call: Dict[str, Any], sessions: Dict[str, str]) -> str: def close(self, sample: Sample) -> None: """Guaranteed teardown (LIN-533): end every open tool session for this trajectory. - The engine calls this from ``_run_one``'s ``finally`` on every path — success, crash, and - abort — on the trajectory's own drain thread. Swallows per-session errors so teardown can - never destabilize the drain (``_run_one`` must not raise). A no-op for stateless tools / + The harness calls this from ``ToolAgentHarness.run``'s ``finally`` on every path — success, + crash, and suspension — on the trajectory's own drain thread. Swallows per-session errors + so teardown can never destabilize the drain. A no-op for stateless tools / sessionless trajectories. """ if not self._stateful_tools: diff --git a/unirl/rollout/env/tools/__init__.py b/unirl/rollout/env/tools/__init__.py new file mode 100644 index 000000000..3c1e9a18d --- /dev/null +++ b/unirl/rollout/env/tools/__init__.py @@ -0,0 +1,16 @@ +"""Tools a :class:`~unirl.rollout.env.tool_environment.ToolEnvironment` dispatches to (LIN-492). + +One class per module: the interfaces :class:`~unirl.rollout.env.tools.base.Tool` (stateless) and +:class:`~unirl.rollout.env.tools.base.StatefulTool` (session-scoped), the reference +:class:`~unirl.rollout.env.tools.calculator.CalculatorTool`, the persistent-REPL +:class:`~unirl.rollout.env.tools.sandbox.SandboxTool`, and the deep-research web tools +:class:`~unirl.rollout.env.tools.search.SearchTool` / :class:`~unirl.rollout.env.tools.visit.VisitTool`. +""" + +from unirl.rollout.env.tools.base import StatefulTool, Tool +from unirl.rollout.env.tools.calculator import CalculatorTool +from unirl.rollout.env.tools.sandbox import SandboxTool +from unirl.rollout.env.tools.search import SearchTool +from unirl.rollout.env.tools.visit import VisitTool + +__all__ = ["Tool", "StatefulTool", "CalculatorTool", "SandboxTool", "SearchTool", "VisitTool"] diff --git a/unirl/rollout/loop/tools/tool.py b/unirl/rollout/env/tools/base.py similarity index 88% rename from unirl/rollout/loop/tools/tool.py rename to unirl/rollout/env/tools/base.py index 887ff3cba..cfa1fd8e0 100644 --- a/unirl/rollout/loop/tools/tool.py +++ b/unirl/rollout/env/tools/base.py @@ -3,7 +3,7 @@ A ``Tool`` is the two things the environment needs: a JSON schema (so the rollout prompt can advertise the tool to the model via ``tokenizer.apply_chat_template(tools=...)``) and an executor (run the parsed call, return a text result). One concrete tool per module; see -:class:`~unirl.rollout.loop.tools.calculator.CalculatorTool`. +:class:`~unirl.rollout.env.tools.calculator.CalculatorTool`. """ from __future__ import annotations @@ -28,7 +28,7 @@ def json_schema(self) -> Dict[str, Any]: def execute(self, arguments: Dict[str, Any]) -> str: """Run the tool on parsed ``arguments`` and return the result as text. - May raise on bad input — :class:`~unirl.rollout.loop.tool_environment.ToolEnvironment` + May raise on bad input — :class:`~unirl.rollout.env.tool_environment.ToolEnvironment` catches and surfaces the error to the model as the observation, so the policy can recover. """ ... @@ -39,7 +39,7 @@ class StatefulTool(Tool): Where :class:`Tool` is a pure function (args in, text out, holds nothing), a ``StatefulTool`` carries state across turns — a code-interpreter namespace, an editing canvas, a connection. - :class:`~unirl.rollout.loop.tool_environment.ToolEnvironment` dispatches on + :class:`~unirl.rollout.env.tool_environment.ToolEnvironment` dispatches on ``isinstance(tool, StatefulTool)`` (one protocol, no code fork), so the stateless ``Tool`` path is byte-for-byte unchanged. @@ -52,8 +52,8 @@ class StatefulTool(Tool): ``execute_session``, which runs off-loop in an executor. - ``execute_session(session_id, arguments)`` — per turn. Operates on the (lazily opened) per-session handle; runs in the executor via - :meth:`~unirl.rollout.loop.tool_environment.ToolEnvironment.step`. - - ``session_end(session_id)`` — once, guaranteed: the engine's ``finally`` hook calls it even on + :meth:`~unirl.rollout.env.tool_environment.ToolEnvironment.step`. + - ``session_end(session_id)`` — once, guaranteed: the harness's ``finally`` hook calls it even on a crashed/aborted trajectory (via ``ToolEnvironment.close``). Must be **idempotent**, a no-op on an unknown/never-opened id, and **must not raise**. @@ -69,7 +69,7 @@ def session_start(self, session_id: str, context: Dict[str, Any]) -> None: def execute_session(self, session_id: str, arguments: Dict[str, Any]) -> str: """Run the tool for ``session_id`` on parsed ``arguments``; return the result as text. - May raise on bad input — :class:`~unirl.rollout.loop.tool_environment.ToolEnvironment` + May raise on bad input — :class:`~unirl.rollout.env.tool_environment.ToolEnvironment` catches and surfaces the error to the model as the observation. """ ... diff --git a/unirl/rollout/loop/tools/calculator.py b/unirl/rollout/env/tools/calculator.py similarity index 96% rename from unirl/rollout/loop/tools/calculator.py rename to unirl/rollout/env/tools/calculator.py index a31ca196a..96b12458d 100644 --- a/unirl/rollout/loop/tools/calculator.py +++ b/unirl/rollout/env/tools/calculator.py @@ -1,6 +1,6 @@ """CalculatorTool — a safe arithmetic tool (LIN-492). -The first concrete :class:`~unirl.rollout.loop.tools.tool.Tool`: evaluates an arithmetic +The first concrete :class:`~unirl.rollout.env.tools.base.Tool`: evaluates an arithmetic expression with no infrastructure (no sandbox, no network). Evaluation is a hand-rolled AST walk over a numeric whitelist — never :func:`eval` — so a hallucinated ``__import__(...)``, name, or attribute access is rejected rather than executed. @@ -12,7 +12,7 @@ import operator from typing import Any, Dict -from unirl.rollout.loop.tools.tool import Tool +from unirl.rollout.env.tools.base import Tool # Binary / unary operators the calculator allows. Anything else (names, calls, attributes, # comparisons, subscripts, ...) is rejected by ``_eval`` below. diff --git a/unirl/rollout/loop/tools/sandbox.py b/unirl/rollout/env/tools/sandbox.py similarity index 97% rename from unirl/rollout/loop/tools/sandbox.py rename to unirl/rollout/env/tools/sandbox.py index a7e03c5d6..850e1af97 100644 --- a/unirl/rollout/loop/tools/sandbox.py +++ b/unirl/rollout/env/tools/sandbox.py @@ -1,9 +1,9 @@ """SandboxTool — a persistent Python-REPL subprocess tool (LIN-533). -The first out-of-process :class:`~unirl.rollout.loop.tools.tool.StatefulTool`: each session owns a +The first out-of-process :class:`~unirl.rollout.env.tools.base.StatefulTool`: each session owns a long-lived ``python`` subprocess whose global namespace **persists across turns**, so ``x = 40`` on one turn and ``x + 2`` on the next returns ``42`` — the cross-turn capability a stateless -:class:`~unirl.rollout.loop.tools.tool.Tool` cannot express. +:class:`~unirl.rollout.env.tools.base.Tool` cannot express. Lifecycle (all off the shared event loop — ``execute_session``/``session_end`` run in ``ToolEnvironment``'s executor): @@ -31,7 +31,7 @@ import threading from typing import Any, Dict, Optional -from unirl.rollout.loop.tools.tool import StatefulTool +from unirl.rollout.env.tools.base import StatefulTool # Runs in the child process. Protocol: read one JSON line ``{"code": ...}`` from stdin; exec it in a # persistent namespace with stdout+stderr captured; if the last top-level statement is an expression, diff --git a/unirl/rollout/loop/tools/search.py b/unirl/rollout/env/tools/search.py similarity index 97% rename from unirl/rollout/loop/tools/search.py rename to unirl/rollout/env/tools/search.py index 0462bd96d..906de42e4 100644 --- a/unirl/rollout/loop/tools/search.py +++ b/unirl/rollout/env/tools/search.py @@ -1,6 +1,6 @@ """SearchTool — batched web search via Serper or SerpApi (LIN-519, hardened). -A concrete :class:`~unirl.rollout.loop.tools.tool.Tool` for the deep-research +A concrete :class:`~unirl.rollout.env.tools.base.Tool` for the deep-research agent: given an array of query strings it returns the top web results per query as text. Two providers, selected by ``$SEARCH_PROVIDER`` (or the constructor): @@ -20,7 +20,7 @@ import requests -from unirl.rollout.loop.tools.tool import Tool +from unirl.rollout.env.tools.base import Tool _SERPER_URL = "https://google.serper.dev/search" _SERPAPI_URL = "https://serpapi.com/search" diff --git a/unirl/rollout/loop/tools/visit.py b/unirl/rollout/env/tools/visit.py similarity index 98% rename from unirl/rollout/loop/tools/visit.py rename to unirl/rollout/env/tools/visit.py index d0fb17c08..b67d9f813 100644 --- a/unirl/rollout/loop/tools/visit.py +++ b/unirl/rollout/env/tools/visit.py @@ -1,6 +1,6 @@ """VisitTool — read webpage(s) and summarize toward a goal (LIN-519, hardened). -A concrete :class:`~unirl.rollout.loop.tools.tool.Tool` for the deep-research +A concrete :class:`~unirl.rollout.env.tools.base.Tool` for the deep-research agent: fetch a URL's content with the Jina reader (needs ``$JINA_API_KEYS``) and summarize the parts relevant to a stated goal with an OpenAI-compatible LLM (hosted out-of-band; ``$SUMMARY_URL`` / ``$SUMMARY_MODEL`` or the constructor @@ -24,7 +24,7 @@ import requests -from unirl.rollout.loop.tools.tool import Tool +from unirl.rollout.env.tools.base import Tool _JINA_READ = "https://r.jina.ai/" # Structured extractor (mirrors AReaL's EXTRACTOR_PROMPT): evidence + summary. diff --git a/unirl/rollout/harness/__init__.py b/unirl/rollout/harness/__init__.py new file mode 100644 index 000000000..067055696 --- /dev/null +++ b/unirl/rollout/harness/__init__.py @@ -0,0 +1,17 @@ +"""Rollout harnesses: worker-side task-internal control flow, as plugins. + +A harness owns the sequence dictated by TASK SEMANTICS — how many model/tool +turns, what ends an episode, how multi-stage flows chain — the third kind of +sequence owner next to training policy (trainer-side step loops) and engine +contracts (driver-side pumps). It runs INSIDE the rollout worker, so +multi-turn env state and intermediate tensors never cross to the driver. + +``protocol.py`` holds the narrow boundary (``RolloutHarness`` / +``HarnessContext`` / ``HarnessOutcome``); one harness = one module beside it. +Harnesses share the boundary, never a loop — different tasks are different +control flows, and forcing them through one shared loop is false unification. + +Deliberately empty otherwise: harness modules must stay ray/torch-free to +import (they are exercised by CPU harnesses with fake engines/envs), so this +init imports nothing and consumers import the submodules directly. +""" diff --git a/unirl/rollout/harness/protocol.py b/unirl/rollout/harness/protocol.py new file mode 100644 index 000000000..69c25b536 --- /dev/null +++ b/unirl/rollout/harness/protocol.py @@ -0,0 +1,92 @@ +"""The harness boundary: what a task-internal control flow may see and return. + +Kept deliberately narrow. The vocabulary test is the guard: nothing here may +speak trainer words (``rollout_id`` / ``sync_interval`` / rewards scoring / +GRPO grouping / checkpoint cadence) or deployment words (wake/sleep/offload — +engine residency is the hosting runtime's job, applied at ``generate`` +boundaries per deployment, so the same harness runs colocate, separate, or +trainside unchanged). +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Callable, Literal, Mapping, Optional, Protocol + +if TYPE_CHECKING: + from unirl.types.sample import Sample + + +@dataclass(frozen=True) +class HarnessOutcome: + """What one task run produced. + + ``completed`` — the trajectory is terminal; ``sample`` is the full trace. + ``suspended`` — cooperatively stopped at a harness-chosen safe point; + ``sample`` is the checkpoint. Whether it can be RESUBMITTED depends on + the environment, not this contract: a stateless env (tools derive state + from the Sample) resumes; a stateful env (own episodes/sessions, torn + down by ``close``) cannot, and its recipes must pair with + ``tail_policy: drop``. Config-time validation of that pairing is + deferred to the config-selected-harness step. + ``failed`` — the task hit an infrastructure fault (backend outage, tool + timeout, context overflow); ``sample`` is the partial trace. The hosting + runtime marks it so a fault never enters advantage math as a legitimate + low-scoring sibling. + + ``env_reward`` — an env-sourced per-trajectory return, when the task's + environment emits one (``info["reward"]``; last value wins). Captured here + because reading it is task semantics; ATTACHING it to the trace is tensor + work the hosting runtime does, keeping harness modules torch-free. + """ + + sample: "Sample" + status: Literal["completed", "suspended", "failed"] + env_reward: Optional[float] = None + + +@dataclass(frozen=True) +class HarnessContext: + """The runtime surface a harness runs against: named engines + a stop probe. + + Engine names ("policy", "ar", "diffusion", ...) are the hosting runtime's + local dependency names from its config — not a global registry; the + harness file states which names it uses, so the call graph stays readable. + """ + + engines: Mapping[str, Callable[["Sample"], "Sample"]] = field(default_factory=dict) + suspend: Callable[[], bool] = lambda: False + + def generate(self, engine: str, sample: "Sample") -> "Sample": + """One blocking model call on the named engine.""" + try: + gen = self.engines[engine] + except KeyError: + raise KeyError( + f"harness asked for engine {engine!r}; this runtime provides {sorted(self.engines)}" + ) from None + return gen(sample) + + def suspend_requested(self) -> bool: + """True once the runtime wants a cooperative stop (quiesce / weight sync). + + The harness checks this at ITS safe points — between turns, between + stages — and returns a ``suspended`` outcome; the runtime never + interrupts mid-generation. + """ + return self.suspend() + + +class RolloutHarness(Protocol): + """One task's internal control flow, hosted by a rollout-worker runtime. + + ``run`` must not raise for task-level faults — catch and return a + ``failed`` outcome so the partial trace survives; the hosting runtime + keeps a last-resort net for harness bugs. Instances are shared across + worker threads: keep all per-task state local to ``run``. + """ + + def run(self, request: "Sample", context: HarnessContext) -> HarnessOutcome: ... + + +__all__ = ["HarnessContext", "HarnessOutcome", "RolloutHarness"] diff --git a/unirl/rollout/harness/tool_agent.py b/unirl/rollout/harness/tool_agent.py new file mode 100644 index 000000000..2c586c174 --- /dev/null +++ b/unirl/rollout/harness/tool_agent.py @@ -0,0 +1,95 @@ +"""ToolAgentHarness — the environment-driven multi-turn agent loop (LIN-492/531). + +The task-semantics half of what ``AgenticRolloutEngine._run_one`` used to +inline (and the successor of the deleted ``AgentLoop`` prototype): each +turn forks a one-sample continuation, the ``"policy"`` engine fills it, and +the ENVIRONMENT decides what happens next — it parses the model's output +(e.g. a tool call), returns an observation that re-enters the chain as a +mask-0 input Part, and signals ``done``. The loop holds no other control +decision; queueing, concurrency, buffers, and abort belong to the hosting +runtime. + +Resume-aware: ``turns_done = len(request.gen_parts())``, so a carried partial +continues from where it stopped (``env.reset`` is idempotent/turn-derived). +Suspension is checked at the top of each turn — the in-flight turn always +finishes naturally first (turn boundary). Whether a suspended trajectory can +actually RESUME is the env's property, not this loop's: stateless tool envs +re-derive state from the Sample; stateful envs (ALFWorld episodes, persistent +sessions) are torn down by ``close`` on suspension, so their recipes pair with +``tail_policy: drop``. +""" + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING, Any, Optional + +from unirl.rollout.harness.protocol import HarnessContext, HarnessOutcome + +if TYPE_CHECKING: + from unirl.rollout.env.protocol import Environment + from unirl.types.sample import Sample + +logger = logging.getLogger(__name__) + + +class ToolAgentHarness: + """``generate -> env.step -> observe`` until ``done`` / ``max_turns``, on the ``"policy"`` engine. + + ``env`` must be re-entrant (one shared instance serves concurrent + trajectories on their own threads); ``sampling`` is only READ each turn + (``fork`` builds a fresh gen Part per call), so sharing this harness + across worker threads is safe as long as ``env.step`` is. + """ + + ENGINE = "policy" + + def __init__(self, *, env: "Environment", sampling: Any, max_turns: int) -> None: + self.env = env + self.sampling = sampling + self.max_turns = int(max_turns) + + def run(self, request: "Sample", context: HarnessContext) -> HarnessOutcome: + sample = request + env_reward: Optional[float] = None + try: + sample = self.env.reset(request) # [input(1)], root id = prompt id + turns_done = len(sample.gen_parts()) + for _ in range(self.max_turns - turns_done): + if context.suspend_requested(): # partial rollout: checkpoint at the turn boundary + return HarnessOutcome(sample, "suspended", env_reward) + sample = context.generate(self.ENGINE, sample.fork(1, sampling_params=self.sampling)) # +[gen(1)] + observation, done, info = self.env.step(sample) # blocking tool boundary, own thread + # Env-sourced reward (LIN-519): interactive envs (ALFWorld, …) return a + # per-trajectory return in ``info["reward"]`` (last value = the episode + # return); tool-only envs (calculator/search) omit it — a no-op here. + if isinstance(info, dict) and info.get("reward") is not None: + env_reward = float(info["reward"]) + if done: + return HarnessOutcome(sample, "completed", env_reward) + if observation is not None: + sample = sample.observe(observation) # +[obs(1)] + return HarnessOutcome(sample, "completed", env_reward) # max_turns reached = terminal + except Exception as exc: # noqa: BLE001 — isolate: one bad trajectory must not sink the drain + # Task-level fault (backend outage, tool timeout, context overflow): + # return the partial trace as ``failed`` — the runtime marks it so a + # fault never enters advantage math as a legitimate low-scoring + # sibling. Any partial ``env_reward`` is deliberately dropped: a + # reward collected before the fault does not describe a complete + # trajectory. + logger.warning("ToolAgentHarness: trajectory failed: %s", exc, exc_info=True) + return HarnessOutcome(sample, "failed") + finally: + # Guaranteed teardown (LIN-533): end any open tool sessions / episodes + # for this trajectory — on success, crash, AND suspension. Duck-typed + # like ``tool_schemas`` so envs without ``close`` are unaffected, and + # wrapped so a teardown error can never re-raise. + close = getattr(self.env, "close", None) + if close is not None: + try: + close(sample) + except Exception: # noqa: BLE001 — teardown must not sink the drain + logger.warning("ToolAgentHarness: env.close failed during teardown", exc_info=True) + + +__all__ = ["ToolAgentHarness"] diff --git a/unirl/rollout/loop/__init__.py b/unirl/rollout/loop/__init__.py deleted file mode 100644 index 2f8625225..000000000 --- a/unirl/rollout/loop/__init__.py +++ /dev/null @@ -1,25 +0,0 @@ -"""Driver-side agent-loop package (LIN-492). See ``unirl/rollout/loop/README.md``. - -One class per module: :class:`~unirl.rollout.loop.engine_port.RolloutEnginePort`, -:class:`~unirl.rollout.loop.environment.Environment`, -:class:`~unirl.rollout.loop.agent_loop.AgentLoop`, and the first concrete environment -:class:`~unirl.rollout.loop.tool_environment.ToolEnvironment` (with its -:mod:`~unirl.rollout.loop.tools`). -""" - -from unirl.rollout.loop.agent_loop import AgentLoop -from unirl.rollout.loop.engine_port import RolloutEnginePort -from unirl.rollout.loop.environment import Environment -from unirl.rollout.loop.tool_environment import ToolEnvironment, parse_tool_call -from unirl.rollout.loop.tools import CalculatorTool, StatefulTool, Tool - -__all__ = [ - "AgentLoop", - "Environment", - "RolloutEnginePort", - "ToolEnvironment", - "parse_tool_call", - "Tool", - "StatefulTool", - "CalculatorTool", -] diff --git a/unirl/rollout/loop/agent_loop.py b/unirl/rollout/loop/agent_loop.py deleted file mode 100644 index a20638da8..000000000 --- a/unirl/rollout/loop/agent_loop.py +++ /dev/null @@ -1,65 +0,0 @@ -"""AgentLoop — a synchronous, environment-driven multi-turn rollout driver (LIN-492). - -See ``unirl/rollout/loop/README.md``. One generic loop runs multi-turn rollout over the existing -synchronous rollout engine, threading state through the ``Sample``/``Part`` model. Each turn -forks a continuation, the engine fills it, and the **environment** decides what happens next: it -consumes the model's output (e.g. parses a tool call), returns an observation that re-enters the -chain as a mask-0 input Part, and signals ``done`` when the trajectory is complete (e.g. the model -stopped emitting tool calls). The loop itself holds **no control decision** — it is purely -mechanical: ``generate -> env.step -> observe -> repeat``, bounded by ``max_turns``. - -This matches verl / slime / areal / relax: the agent loop is driven by the model's tool call, not -a fixed plan. Fixed multi-stage pipelines (e.g. the AR→diffusion ``ComposedRolloutEngine``) are -*not* agent loops and do not use this class. - -Prototype scope: synchronous only (the engine's ``generate`` is a blocking call). -""" - -from __future__ import annotations - -from unirl.rollout.loop.engine_port import RolloutEnginePort -from unirl.rollout.loop.environment import Environment -from unirl.types.sample import Sample -from unirl.types.sampling import BaseSamplingParams, total_samples_per_prompt - - -class AgentLoop: - """One generic, SYNCHRONOUS, environment-driven multi-turn rollout loop. - - The loop owns only mechanical control flow — not generation (engine), world dynamics or - termination (environment), scoring, or training. Behaviour is set by: - - - ``environment`` — **the driver** (required). Each turn its :meth:`Environment.step` consumes - the model's latest output and returns ``(observation, done, info)``; ``done`` ends the - trajectory (the tool-call-driven stop condition). ``observation`` re-enters as a mask-0 Part. - - ``sampling_params`` — one fixed config used every turn. The GRPO fan-out ``n`` is read from it - for turn 0 (:func:`total_samples_per_prompt`); continuations fork one sample each. - - ``max_turns`` — hard safety bound on turns. - """ - - def __init__( - self, - environment: Environment, - sampling_params: BaseSamplingParams, - max_turns: int = 8, - ) -> None: - self.environment = environment - self.sampling_params = sampling_params - self.max_turns = max_turns - - def run(self, engine: RolloutEnginePort, request: Sample) -> Sample: - """Drive the episode: ``fork -> generate -> env.step -> observe`` until ``done`` / ``max_turns``.""" - sample = self.environment.reset(request) - branch = total_samples_per_prompt(self.sampling_params) # GRPO fan-out on the first turn - for _ in range(self.max_turns): - sample = engine.generate(sample.fork(branch, sampling_params=self.sampling_params)) - observation, done, _info = self.environment.step(sample) # the environment decides termination - if done: - break - if observation is not None: - sample = sample.observe(observation) - branch = 1 # continuations are one sample each - return sample - - -__all__ = ["AgentLoop"] diff --git a/unirl/rollout/loop/engine_port.py b/unirl/rollout/loop/engine_port.py deleted file mode 100644 index 6caf42385..000000000 --- a/unirl/rollout/loop/engine_port.py +++ /dev/null @@ -1,26 +0,0 @@ -"""RolloutEnginePort — the generation seam the agent loop calls (LIN-492). - -See ``unirl/rollout/loop/README.md``. A structural ``Protocol`` for a single-turn -engine; ``SyncRolloutEngine`` is its nominal runtime counterpart. -""" - -from __future__ import annotations - -from typing import Protocol - -from unirl.types.sample import Sample - - -class RolloutEnginePort(Protocol): - """What the loop calls to generate one model turn.""" - - def generate(self, sample: Sample) -> Sample: - """Fill the request Sample's frontier gen Part and return it. - - One turn is always one ``Sample`` (never a trajectory list) — that - list belongs to the coordinator contract. - """ - ... - - -__all__ = ["RolloutEnginePort"] diff --git a/unirl/rollout/loop/tools/__init__.py b/unirl/rollout/loop/tools/__init__.py deleted file mode 100644 index 28d227333..000000000 --- a/unirl/rollout/loop/tools/__init__.py +++ /dev/null @@ -1,16 +0,0 @@ -"""Tools a :class:`~unirl.rollout.loop.tool_environment.ToolEnvironment` dispatches to (LIN-492). - -One class per module: the interfaces :class:`~unirl.rollout.loop.tools.tool.Tool` (stateless) and -:class:`~unirl.rollout.loop.tools.tool.StatefulTool` (session-scoped), the reference -:class:`~unirl.rollout.loop.tools.calculator.CalculatorTool`, the persistent-REPL -:class:`~unirl.rollout.loop.tools.sandbox.SandboxTool`, and the deep-research web tools -:class:`~unirl.rollout.loop.tools.search.SearchTool` / :class:`~unirl.rollout.loop.tools.visit.VisitTool`. -""" - -from unirl.rollout.loop.tools.calculator import CalculatorTool -from unirl.rollout.loop.tools.sandbox import SandboxTool -from unirl.rollout.loop.tools.search import SearchTool -from unirl.rollout.loop.tools.tool import StatefulTool, Tool -from unirl.rollout.loop.tools.visit import VisitTool - -__all__ = ["Tool", "StatefulTool", "CalculatorTool", "SandboxTool", "SearchTool", "VisitTool"] diff --git a/unirl/train_agentic_env.py b/unirl/train_agentic_env.py index b0d32ea58..5ff95cd0c 100755 --- a/unirl/train_agentic_env.py +++ b/unirl/train_agentic_env.py @@ -3,7 +3,7 @@ Drives :class:`unirl.trainer.agentic_env.AgenticEnvTrainer` over the :class:`~unirl.rollout.engine.agentic.engine.AgenticRolloutEngine` with an interactive -:class:`~unirl.rollout.loop.environment.Environment`. Sibling of ``train_agentic.py``; +:class:`~unirl.rollout.env.protocol.Environment`. Sibling of ``train_agentic.py``; the reward is the environment's own per-trajectory return — task-success or a shaped signal, attached to each trajectory by the engine — so no reward backend is scored and the recipe's ``reward`` block is built but unused. @@ -14,7 +14,7 @@ python -m unirl.train_agentic_env --config-name=alfworld/alfworld_grpo num_devices=8 This entrypoint serves every env-reward agentic recipe; ALFWorld -(:class:`~unirl.rollout.loop.alfworld_env.AlfworldEnv`, ``examples/alfworld/``) is the +(:class:`~unirl.rollout.env.alfworld.AlfworldEnv`, ``examples/alfworld/``) is the reference environment. """ diff --git a/unirl/trainer/async_ar.py b/unirl/trainer/async_ar.py index 5bda5f24b..756fd41e5 100644 --- a/unirl/trainer/async_ar.py +++ b/unirl/trainer/async_ar.py @@ -45,7 +45,7 @@ from unirl.distributed.group.placement import placement, remote from unirl.distributed.tensor import hydrate from unirl.models.qwen3_5.validation import validate_qwen3_5_training_contract -from unirl.rollout.engine.asynchronous import AsyncBatchRolloutEngine +from unirl.rollout.engine.asynchronous import AsyncBatchRolloutEngine, launch_ceiling from unirl.train.stack import TrainStepResult from unirl.trainer.ar import ARTrainer from unirl.trainer.base import BaseTrainer, build_sampling_dict @@ -422,7 +422,7 @@ def _next_step( """ engine = self._async_engine while True: - ceiling = min(num_rollouts, ((rollout_id // interval) + 1 + stale) * interval) + ceiling = launch_ceiling(rollout_id, sync_interval=interval, max_staleness=stale, num_rollouts=num_rollouts) while engine.next_gen_id < ceiling and engine.inflight < M: engine.submit(self._build_async_sample(engine.next_gen_id)) engine.poll() diff --git a/unirl/trainer/async_diffusion.py b/unirl/trainer/async_diffusion.py index 83ac178fe..99d02a525 100644 --- a/unirl/trainer/async_diffusion.py +++ b/unirl/trainer/async_diffusion.py @@ -50,7 +50,7 @@ import torch from unirl.distributed.tensor import hydrate -from unirl.rollout.engine.asynchronous import AsyncBatchRolloutEngine +from unirl.rollout.engine.asynchronous import AsyncBatchRolloutEngine, launch_ceiling from unirl.train.stack import TrainStepResult from unirl.trainer.diffusion import DiffusionTrainer from unirl.types.sample import Sample @@ -264,7 +264,7 @@ def _next_step( """ engine = self._async_engine while True: - ceiling = min(num_rollouts, ((rollout_id // interval) + 1 + stale) * interval) + ceiling = launch_ceiling(rollout_id, sync_interval=interval, max_staleness=stale, num_rollouts=num_rollouts) engine.poll() while engine.next_gen_id < ceiling and engine.inflight < M: engine.submit(self._build_async_sample(engine.next_gen_id)) diff --git a/unirl/types/sample.py b/unirl/types/sample.py index 797207d22..7af8a4394 100644 --- a/unirl/types/sample.py +++ b/unirl/types/sample.py @@ -786,7 +786,7 @@ def fork( def observe(self, observation: Primitive, *, role: str = "tool") -> "Sample": """Append an observation as a branch-1, mask-0 *input* Part off the frontier. - The world-response half of an agentic turn (``unirl/rollout/loop/README.md``): the + The world-response half of an agentic turn (``unirl/rollout/env/README.md``): the observation rides as a chained input Part — one child per frontier sample, ids extended by ``/0`` — carrying no ``sampling_params``. So it is excluded from :meth:`gen_parts` (never trained) and surfaced to the next turn by diff --git a/unirl/utils/prepare_alfworld.py b/unirl/utils/prepare_alfworld.py index daa8f0569..f06ae7239 100644 --- a/unirl/utils/prepare_alfworld.py +++ b/unirl/utils/prepare_alfworld.py @@ -5,7 +5,7 @@ ``{"prompt": , "metadata": {"game_index": i}}``, so ``MultimodalRLDataSource`` drives N rollouts and :meth:`AlfworldEnv.reset` maps each ``game_index`` to a game. The index ordering matches -:func:`unirl.rollout.loop.alfworld_env.list_alfworld_games`, so the row and the env +:func:`unirl.rollout.env.alfworld.list_alfworld_games`, so the row and the env agree on which game an index selects (and the ``n`` GRPO siblings of a prompt share the same game). The prompt is a non-empty placeholder only because the data source rejects empty-prompt rows — :meth:`AlfworldEnv.reset` discards it and uses the env observation. @@ -19,7 +19,7 @@ import json import os -from unirl.rollout.loop.alfworld_env import list_alfworld_games +from unirl.rollout.env.alfworld import list_alfworld_games # Non-empty so the data source keeps the row; AlfworldEnv.reset() replaces it with the # environment's initial observation, so the text itself is never used for generation.