refactor(rollout): dedup async driver plumbing; extract ToolAgentHarness

.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/

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 ⚡

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

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

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

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:

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:

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:

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.

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 ──