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.