From ba67d66ee3f9f41019af4b07b811f0509bd88861 Mon Sep 17 00:00:00 2001
From: FlappyBear <pty12345@mail.ustc.edu.cn>
Date: Mon, 21 Sep 2026 16:30:11 +0800
Subject: [PATCH 1/2] feat(rollout): collect complete greedy episodes and
 paired baselines

---
 agent_r1/agent_flow/agent_env_loop.py         |  15 +-
 agent_r1/agent_flow/agent_flow.py             | 180 +++++-
 agent_r1/agent_flow/rollout_utils.py          | 166 +++++
 agent_r1/agent_flow/single_step_agent_flow.py |   5 +-
 agent_r1/env/envs/tool.py                     |   2 +-
 docs/tutorials/remax-greedy-rollout.md        |  91 +++
 docs/zh/tutorials/remax-greedy-rollout.md     |  84 +++
 mkdocs.yml                                    |   2 +
 recipes/alfworld/alfworld_agent_flow.py       |  13 +-
 recipes/alfworld/env/alfworld_wrapper.py      |   5 +-
 recipes/hotpotqa/hotpotqa_agent_flow.py       |   8 +-
 .../paper_search/paper_search_agent_flow.py   |   8 +-
 recipes/webshop/webshop_agent_flow.py         |  12 +-
 tests/test_greedy_rollouts.py                 | 593 ++++++++++++++++++
 tests/test_greedy_tensor_contract.py          |  75 +++
 15 files changed, 1220 insertions(+), 39 deletions(-)
 create mode 100644 agent_r1/agent_flow/rollout_utils.py
 create mode 100644 docs/tutorials/remax-greedy-rollout.md
 create mode 100644 docs/zh/tutorials/remax-greedy-rollout.md
 create mode 100644 tests/test_greedy_rollouts.py
 create mode 100644 tests/test_greedy_tensor_contract.py

diff --git a/agent_r1/agent_flow/agent_env_loop.py b/agent_r1/agent_flow/agent_env_loop.py
index 9aae760..3defbdb 100644
--- a/agent_r1/agent_flow/agent_env_loop.py
+++ b/agent_r1/agent_flow/agent_env_loop.py
@@ -10,6 +10,7 @@
     AgentFlowStep,
     register,
 )
+from agent_r1.agent_flow.rollout_utils import terminal_status
 from agent_r1.env import AgentEnv
 from agent_r1.env.base import Action, Observation
 from verl.utils.profiler import simple_timer
@@ -113,6 +114,7 @@ async def run(self, sampling_params: dict[str, Any], **kwargs) -> AgentFlowOutpu
 
         steps: list = []
         metrics = {}
+        end_status = {"terminated": False, "truncated": True, "termination_reason": "max_steps"}
 
         for step_idx in range(self.max_steps):
             prompt_ids = await self._obs_to_prompt(obs, tools=tools)
@@ -126,6 +128,7 @@ async def run(self, sampling_params: dict[str, Any], **kwargs) -> AgentFlowOutpu
                     self.prompt_length,
                     step_idx,
                 )
+                end_status["termination_reason"] = "prompt_length"
                 break
 
             with simple_timer("generate_sequences", metrics):
@@ -136,6 +139,7 @@ async def run(self, sampling_params: dict[str, Any], **kwargs) -> AgentFlowOutpu
                 )
 
             response_ids = output.token_ids[: self.response_length]
+            generation_info = self._generation_metadata(output)
 
             response_text = await self.loop.run_in_executor(
                 None,
@@ -158,12 +162,21 @@ async def run(self, sampling_params: dict[str, Any], **kwargs) -> AgentFlowOutpu
                     else None
                 ),
                 reward_score=reward,
+                extra_fields=generation_info,
             )
             step = await self._postprocess(step, **kwargs)
             steps.append(step)
 
             if done:
+                info = info or {}
+                env_truncated = bool(info.get("truncated", info.get("TimeLimit.truncated", False)))
+                if env_truncated:
+                    end_status = {"terminated": False, "truncated": True, "termination_reason": "env_truncated"}
+                elif info.get("termination_reason") == "final_answer":
+                    end_status = terminal_status(generation_info)
+                else:
+                    end_status = {"terminated": True, "truncated": False, "termination_reason": "env_done"}
                 break
             obs = next_obs
 
-        return AgentFlowOutput(steps=steps, metrics=metrics)
+        return AgentFlowOutput(steps=steps, metrics=metrics, **end_status)
diff --git a/agent_r1/agent_flow/agent_flow.py b/agent_r1/agent_flow/agent_flow.py
index 9281efa..16a26db 100644
--- a/agent_r1/agent_flow/agent_flow.py
+++ b/agent_r1/agent_flow/agent_flow.py
@@ -15,6 +15,7 @@
 import logging
 import os
 from abc import ABC, abstractmethod
+from copy import deepcopy
 from typing import Any, Optional
 from uuid import uuid4
 
@@ -28,6 +29,14 @@
 from tensordict import TensorDict
 from transformers import AutoProcessor, AutoTokenizer
 
+from agent_r1.agent_flow.rollout_utils import (
+    ReMaxRolloutCollection,
+    build_sampling_params,
+    generation_metadata,
+    normalize_source_uid,
+    pair_baseline_rewards,
+    summarize_greedy_baselines,
+)
 from agent_r1.reward_loop.reward_loop import RewardLoopWorker
 from verl.experimental.agent_loop.agent_loop import (
     AsyncLLMServerManager,
@@ -124,6 +133,13 @@ class AgentFlowOutput(BaseModel):
     """List of agent flow steps."""
     metrics: AgentFlowMetrics
     """Auxiliary performance metrics"""
+    source_uid: Optional[str] = None
+    """Original task uid, shared by its sampled and greedy rollouts."""
+    rollout_mode: str = "sample"
+    terminated: bool = False
+    truncated: bool = False
+    termination_reason: str = "unknown"
+    """Unknown is retained for custom flows which have not declared an end status."""
 
 
 class AgentFlowBase(ABC):
@@ -163,6 +179,13 @@ def __init__(
         self.system_prompt = initialize_system_prompt(self.tokenizer, **self.apply_chat_template_kwargs)
         self.loop = get_event_loop()
 
+    def _generation_metadata(self, output) -> dict:
+        return generation_metadata(
+            output,
+            self.config.actor_rollout_ref.rollout.response_length,
+            getattr(self.tokenizer, "eos_token_id", None),
+        )
+
     async def process_vision_info(self, messages: list[dict]) -> dict:
         """Extract images and videos from messages.
 
@@ -545,17 +568,10 @@ async def generate_sequences(self, batch: DataProto) -> DataProto:
             response_mask: | 1, 1, 1, ..., 1, 1 | 0, 0, .., 0, 0 | 1, 1, 1, ..., 1, 1 | 0, 0, ..., 0|
         """
         config = self.config.actor_rollout_ref.rollout
-        sampling_params = dict(
-            temperature=config.temperature,
-            top_p=config.top_p,
-            repetition_penalty=1.0,
-            logprobs=config.calculate_log_probs,
-        )
-
-        # override sampling params for validation
-        if batch.meta_info.get("validate", False):
-            sampling_params["top_p"] = config.val_kwargs.top_p
-            sampling_params["temperature"] = config.val_kwargs.temperature
+        sampling_params = build_sampling_params(config, batch.meta_info)
+        rollout_mode = batch.meta_info.get("rollout_mode", "sample")
+        if "uid" not in batch.non_tensor_batch:
+            batch.non_tensor_batch["uid"] = np.array([uuid4().hex for _ in range(len(batch))], dtype=object)
 
         # by default, we assume it's a single turn agent
         if "agent_name" not in batch.non_tensor_batch:
@@ -593,10 +609,13 @@ async def generate_sequences(self, batch: DataProto) -> DataProto:
             kwargs = {k: v[i] for k, v in batch.non_tensor_batch.items()}
             tasks.append(
                 asyncio.create_task(
-                    self._run_agent_flow(sampling_params, trajectory_info[i], trace=trace_this_sample, **kwargs)
+                    self._run_agent_flow(dict(sampling_params), trajectory_info[i], trace=trace_this_sample, **kwargs)
                 )
             )
         outputs = await asyncio.gather(*tasks)
+        for i, output in enumerate(outputs):
+            output.source_uid = normalize_source_uid(batch.non_tensor_batch["uid"][i])
+            output.rollout_mode = rollout_mode
 
         output = self._postprocess(outputs)
         return output
@@ -639,6 +658,8 @@ async def _run_agent_flow(
 
     def _postprocess(self, inputs: list[AgentFlowOutput]) -> DataProto:
         """Process the padded outputs from _run_agent_flow and combine them into a batch."""
+        if not inputs:
+            raise ValueError("Cannot postprocess an empty rollout batch")
         num_steps = []
         trajectory_uids = []
         step_indices = []
@@ -654,11 +675,26 @@ def _postprocess(self, inputs: list[AgentFlowOutput]) -> DataProto:
         reward_tensors = []
         response_logprobs_list = []
         routed_experts_list = []
+        source_uids = []
+        rollout_modes = []
+        terminated = []
+        truncated = []
+        termination_reasons = []
         for input in inputs:
             num_step = len(input.steps)
+            if not num_step:
+                raise RuntimeError(
+                    f"Rollout for source uid {input.source_uid!r} has no steps "
+                    f"(termination_reason={input.termination_reason!r})"
+                )
             num_steps.append(num_step)
             trajectory_uids.extend([uuid4().hex] * num_step)
             step_indices.extend(range(num_step))
+            source_uids.extend([input.source_uid] * num_step)
+            rollout_modes.extend([input.rollout_mode] * num_step)
+            terminated.extend([False] * (num_step - 1) + [input.terminated])
+            truncated.extend([False] * (num_step - 1) + [input.truncated])
+            termination_reasons.extend(["ongoing"] * (num_step - 1) + [input.termination_reason])
             for step in input.steps:
                 prompt_ids.append(step.prompt_ids)
                 response_ids.append(step.response_ids)
@@ -673,8 +709,10 @@ def _postprocess(self, inputs: list[AgentFlowOutput]) -> DataProto:
                 routed_experts_list.append(step.routed_experts)
                 if step.reward_score is not None:
                     reward_tensor = torch.zeros_like(step.response_mask, dtype=torch.float32)
-                    valid_length = step.response_mask.sum().item()
-                    reward_tensor[0, valid_length - 1] = float(step.reward_score)
+                    valid_positions = step.response_mask[0].nonzero(as_tuple=True)[0]
+                    if not valid_positions.numel():
+                        raise RuntimeError(f"Rollout for source uid {input.source_uid!r} has no response tokens")
+                    reward_tensor[0, valid_positions[-1]] = float(step.reward_score)
                     reward_tensors.append(reward_tensor)
                 else:
                     reward_tensors.append(None)
@@ -712,8 +750,6 @@ def _postprocess(self, inputs: list[AgentFlowOutput]) -> DataProto:
             batch["rm_scores"] = reward_tensor
 
         non_tensor_batch = {
-            "trajectory_uids": np.array(trajectory_uids, dtype=object),
-            "step_indices": np.array(step_indices, dtype=np.int32),
             "__num_turns__": np.array(num_turns, dtype=np.int32),
         }
 
@@ -758,6 +794,17 @@ def _postprocess(self, inputs: list[AgentFlowOutput]) -> DataProto:
             extra_fields[key] = np.array(temp_list, dtype=object)
 
         non_tensor_batch.update(extra_fields)
+        # Provenance/end markers are authoritative and cannot be overwritten by
+        # task-specific extra fields or reward metadata.
+        non_tensor_batch.update(
+            trajectory_uids=np.array(trajectory_uids, dtype=object),
+            step_indices=np.array(step_indices, dtype=np.int32),
+            source_uid=np.array(source_uids, dtype=object),
+            rollout_mode=np.array(rollout_modes, dtype=object),
+            terminated=np.array(terminated, dtype=np.bool_),
+            truncated=np.array(truncated, dtype=np.bool_),
+            termination_reason=np.array(termination_reasons, dtype=object),
+        )
         return DataProto(
             batch=batch,
             non_tensor_batch=non_tensor_batch,
@@ -909,6 +956,72 @@ def _init_agent_flow_workers(self):
                 ).remote(self.config, self.server_handles, self.reward_router_address)
             )
 
+    @staticmethod
+    def _prepare_original_tasks(prompts: DataProto) -> DataProto:
+        """Copy an unrepeated task batch and give each task a stable, unique uid."""
+        if not len(prompts):
+            raise ValueError("A rollout collection needs at least one task")
+        prepared = deepcopy(prompts)
+        uids = prepared.non_tensor_batch.get("uid", [uuid4().hex for _ in range(len(prepared))])
+        uids = [normalize_source_uid(uid) for uid in uids]
+        if len(uids) != len(prepared) or len(set(uids)) != len(uids):
+            raise ValueError("Pass one row per original task, before rollout.n repetition")
+        prepared.non_tensor_batch["uid"] = np.array(uids, dtype=object)
+        return prepared
+
+    def generate_greedy_sequences(self, prompts: DataProto) -> DataProto:
+        """Run one independent, complete greedy episode per original task.
+
+        Uses the same AgentFlow, environment reset, reward function and limits
+        as sampling. Does not change rollout configuration or update the actor.
+        """
+        greedy_batch = self._prepare_original_tasks(prompts)
+        greedy_batch.meta_info["rollout_mode"] = "greedy"
+        return self.generate_sequences(greedy_batch)
+
+    def collect_remax_rollouts(self, prompts: DataProto, num_samples: Optional[int] = None) -> ReMaxRolloutCollection:
+        """Collect paired sample/greedy trajectories, without advantage/update.
+
+        Caller must not update actor weights while this synchronous collection
+        runs. A truncated episode keeps its existing finite-horizon reward; a
+        missing, aborted or unknown-status baseline raises instead of scoring 0.
+        """
+        if prompts.meta_info.get("validate", False):
+            raise ValueError("ReMax collection is a training rollout, not a validation request")
+        if num_samples is None:
+            num_samples = self.config.actor_rollout_ref.rollout.n
+        if isinstance(num_samples, bool) or not isinstance(num_samples, int) or num_samples <= 0:
+            raise ValueError("num_samples must be a positive integer")
+        original = self._prepare_original_tasks(prompts)
+        expected_uids = original.non_tensor_batch["uid"].tolist()
+        greedy = self.generate_greedy_sequences(original)
+        if "rm_scores" not in greedy.batch:
+            raise RuntimeError("Greedy rollout did not return immediate reward scores")
+        step_rewards = (greedy.batch["rm_scores"] * greedy.batch["response_mask"]).sum(dim=-1).tolist()
+        fields = greedy.non_tensor_batch
+        rows = [
+            {
+                "source_uid": fields["source_uid"][i],
+                "trajectory_uid": fields["trajectory_uids"][i],
+                "step_index": int(fields["step_indices"][i]),
+                "rollout_mode": fields["rollout_mode"][i],
+                "terminated": bool(fields["terminated"][i]),
+                "truncated": bool(fields["truncated"][i]),
+                "termination_reason": fields["termination_reason"][i],
+                "reward": reward,
+            }
+            for i, reward in enumerate(step_rewards)
+        ]
+        baselines = summarize_greedy_baselines(rows, expected_uids)
+        sample_batch = original.repeat(repeat_times=num_samples, interleave=True)
+        sample_batch.meta_info["rollout_mode"] = "sample"
+        sampled = self.generate_sequences(sample_batch)
+        paired_rewards = pair_baseline_rewards(sampled.non_tensor_batch["source_uid"].tolist(), baselines)
+        sampled.batch["reward_baselines"] = torch.tensor(
+            paired_rewards, dtype=greedy.batch["rm_scores"].dtype, device=sampled.batch["responses"].device
+        )
+        return ReMaxRolloutCollection(sampled=sampled, greedy=greedy, baselines=baselines)
+
     def generate_sequences(self, prompts: DataProto) -> DataProto:
         """Split input batch and dispatch to agent loop workers.
 
@@ -919,22 +1032,27 @@ def generate_sequences(self, prompts: DataProto) -> DataProto:
             DataProto: Output batch.
         """
 
+        if not len(prompts):
+            raise ValueError("Cannot generate an empty rollout batch")
         self.wake_up()
-        if self.reward_model_manager:
-            self.reward_model_manager.wake_up()
-
-        split_size = (len(prompts) - 1) // len(self.agent_flow_workers) + 1
-        chunks = prompts.split(split_size)
-        outputs = ray.get(
-            [
-                worker.generate_sequences.remote(chunk)
-                for worker, chunk in zip(self.agent_flow_workers, chunks, strict=True)
-            ]
-        )
-        output = DataProto.concat(outputs)
-        self.sleep()
-        if self.reward_model_manager:
-            self.reward_model_manager.sleep()
+        try:
+            if self.reward_model_manager:
+                self.reward_model_manager.wake_up()
+            split_size = (len(prompts) - 1) // len(self.agent_flow_workers) + 1
+            chunks = prompts.split(split_size)
+            outputs = ray.get(
+                [
+                    worker.generate_sequences.remote(chunk)
+                    for worker, chunk in zip(self.agent_flow_workers[: len(chunks)], chunks, strict=True)
+                ]
+            )
+            output = DataProto.concat(outputs)
+        finally:
+            try:
+                self.sleep()
+            finally:
+                if self.reward_model_manager:
+                    self.reward_model_manager.sleep()
 
         # calculate performance metrics
         metrics = [output.meta_info.pop("metrics") for output in outputs]  # List[List[Dict[str, str]]]
diff --git a/agent_r1/agent_flow/rollout_utils.py b/agent_r1/agent_flow/rollout_utils.py
new file mode 100644
index 0000000..9a4da7d
--- /dev/null
+++ b/agent_r1/agent_flow/rollout_utils.py
@@ -0,0 +1,166 @@
+"""Dependency-light greedy rollout controls and trajectory-level reward pairing.
+
+These helpers do not compute advantages or update a policy. Keeping pairing
+independent of tensors makes its invariants testable without a GPU stack.
+"""
+
+import math
+from collections import defaultdict
+from dataclasses import dataclass
+from typing import Any
+
+
+def normalize_source_uid(value: Any) -> str:
+    if value is None or not str(value).strip():
+        raise ValueError("A rollout source uid must be non-empty")
+    return str(value)
+
+
+def build_sampling_params(config: Any, meta_info: dict) -> dict:
+    """Build fresh parameters; greedy overrides validation without changing config.
+
+    ``rollout.n`` is a task repetition count in the trainer, not a backend
+    sampling parameter. Every individual generation request returns one response.
+    """
+    mode = meta_info.get("rollout_mode", "sample")
+    if mode not in ("sample", "greedy"):
+        raise ValueError(f"Unsupported rollout_mode: {mode!r}")
+    if config.response_length <= 0:
+        raise ValueError("rollout.response_length must be positive")
+    params = {
+        "temperature": config.temperature,
+        "top_p": config.top_p,
+        "repetition_penalty": 1.0,
+        "logprobs": config.calculate_log_probs,
+        "max_tokens": config.response_length,
+    }
+    if meta_info.get("validate", False):
+        params.update(temperature=config.val_kwargs.temperature, top_p=config.val_kwargs.top_p)
+    if mode == "greedy":
+        params.update(temperature=0.0, top_p=1.0, logprobs=False)
+    return params
+
+
+def generation_metadata(output: Any, response_length: int, eos_token_id: Any = None) -> dict:
+    """Reject incomplete requests and describe response-budget termination.
+
+    Some verl versions merge backend ``stop`` and ``length`` into ``completed``.
+    Without a raw finish reason, a full-budget response without EOS is marked
+    conservatively as truncated. This is not a guarantee of backend determinism.
+    """
+    stop_reason = getattr(output, "stop_reason", None)
+    finish_reason = getattr(output, "finish_reason", None)
+    if stop_reason in ("abort", "aborted", "error", "failed", "cancelled", "timeout") or finish_reason == "abort":
+        raise RuntimeError(f"Generation did not complete: {finish_reason or stop_reason}")
+    token_ids = output.token_ids
+    if not token_ids:
+        raise RuntimeError("Generation returned an empty response")
+    eos_ids = set(eos_token_id if isinstance(eos_token_id, (list, tuple, set)) else [eos_token_id])
+    at_limit = len(token_ids) >= response_length
+    ends_in_eos = token_ids[-1] in eos_ids
+    truncated = len(token_ids) > response_length or finish_reason == "length" or stop_reason == "length"
+    if finish_reason is None and stop_reason not in ("stop", "length"):
+        truncated = truncated or (at_limit and not ends_in_eos)
+    return {
+        "generation_stop_reason": finish_reason or stop_reason or "unknown",
+        "response_at_limit": at_limit,
+        "response_truncated": truncated,
+    }
+
+
+def terminal_status(metadata: dict, reason: str = "final_answer") -> dict:
+    """Classify a flow which stops after a model response, not an environment done."""
+    truncated = bool(metadata["response_truncated"])
+    return {
+        "terminated": not truncated,
+        "truncated": truncated,
+        "termination_reason": "response_length" if truncated else reason,
+    }
+
+
+@dataclass(frozen=True)
+class GreedyBaseline:
+    source_uid: str
+    trajectory_uid: str
+    reward: float
+    num_steps: int
+    terminated: bool
+    truncated: bool
+    termination_reason: str
+
+
+@dataclass
+class ReMaxRolloutCollection:
+    """Collection only: neither output has advantages or optimizer updates."""
+
+    sampled: Any
+    greedy: Any
+    baselines: dict[str, GreedyBaseline]
+
+
+def summarize_greedy_baselines(rows: list[dict], expected_source_uids: list[str]) -> dict[str, GreedyBaseline]:
+    """Sum immediate rewards over each whole trajectory, independent of row order.
+
+    Require one greedy trajectory per source, contiguous steps and an explicit
+    terminal/truncation marker. Do not invent a zero reward for missing trajectories.
+    """
+    expected = [normalize_source_uid(uid) for uid in expected_source_uids]
+    if not expected or len(set(expected)) != len(expected):
+        raise ValueError("Expected one unique source uid per original task")
+    trajectories = defaultdict(list)
+    for row in rows:
+        if row["rollout_mode"] != "greedy":
+            raise ValueError("Baseline input contains a non-greedy trajectory")
+        trajectories[row["trajectory_uid"]].append(row)
+
+    baselines = {}
+    for trajectory_uid, steps in trajectories.items():
+        steps.sort(key=lambda row: row["step_index"])
+        if [row["step_index"] for row in steps] != list(range(len(steps))):
+            raise ValueError(f"Missing or duplicate steps in trajectory {trajectory_uid}")
+        source_uid = normalize_source_uid(steps[0]["source_uid"])
+        if any(normalize_source_uid(row["source_uid"]) != source_uid for row in steps):
+            raise ValueError(f"Mixed source uids in trajectory {trajectory_uid}")
+        if source_uid in baselines:
+            raise ValueError(f"Multiple greedy trajectories for source uid {source_uid}")
+        if any(row["terminated"] or row["truncated"] for row in steps[:-1]):
+            raise ValueError(f"Trajectory {trajectory_uid} continues after its end marker")
+        last = steps[-1]
+        if bool(last["terminated"]) == bool(last["truncated"]):
+            raise ValueError(f"Trajectory {trajectory_uid} needs exactly one terminal/truncation marker")
+        if last["termination_reason"] in ("unknown", "ongoing", ""):
+            raise ValueError(f"Trajectory {trajectory_uid} has no explicit termination reason")
+        rewards = [float(row["reward"]) for row in steps]
+        if not all(math.isfinite(value) for value in rewards):
+            raise ValueError(f"Non-finite reward in trajectory {trajectory_uid}")
+        try:
+            reward = math.fsum(rewards)
+        except OverflowError as error:
+            raise ValueError(f"Reward overflow in trajectory {trajectory_uid}") from error
+        if not math.isfinite(reward):
+            raise ValueError(f"Non-finite reward in trajectory {trajectory_uid}")
+        baselines[source_uid] = GreedyBaseline(
+            source_uid=source_uid,
+            trajectory_uid=trajectory_uid,
+            reward=reward,
+            num_steps=len(steps),
+            terminated=bool(last["terminated"]),
+            truncated=bool(last["truncated"]),
+            termination_reason=last["termination_reason"],
+        )
+    if set(baselines) != set(expected):
+        missing = set(expected) - set(baselines)
+        unexpected = set(baselines) - set(expected)
+        raise ValueError(f"Greedy/source uid mismatch: missing={sorted(missing)}, unexpected={sorted(unexpected)}")
+    return baselines
+
+
+def pair_baseline_rewards(source_uids: list[str], baselines: dict[str, GreedyBaseline]) -> list[float]:
+    """Broadcast each task's greedy reward to its sampled step rows."""
+    rewards = []
+    for value in source_uids:
+        uid = normalize_source_uid(value)
+        if uid not in baselines:
+            raise ValueError(f"No greedy baseline for sampled source uid {uid}")
+        rewards.append(baselines[uid].reward)
+    return rewards
diff --git a/agent_r1/agent_flow/single_step_agent_flow.py b/agent_r1/agent_flow/single_step_agent_flow.py
index 2746099..b01e45c 100644
--- a/agent_r1/agent_flow/single_step_agent_flow.py
+++ b/agent_r1/agent_flow/single_step_agent_flow.py
@@ -17,6 +17,7 @@
 from uuid import uuid4
 
 from agent_r1.agent_flow.agent_flow import AgentFlowBase, AgentFlowOutput, AgentFlowStep, register
+from agent_r1.agent_flow.rollout_utils import terminal_status
 from verl.utils.profiler import simple_timer
 
 logger = logging.getLogger(__file__)
@@ -60,6 +61,7 @@ async def run(self, sampling_params: dict[str, Any], **kwargs) -> AgentFlowOutpu
             )
 
         response_ids = output.token_ids[: self.response_length]
+        generation_info = self._generation_metadata(output)
 
         step = AgentFlowStep(
             prompt_ids=prompt_ids,
@@ -71,7 +73,8 @@ async def run(self, sampling_params: dict[str, Any], **kwargs) -> AgentFlowOutpu
                 else None
             ),
             multi_modal_data=multi_modal_data,
+            extra_fields=generation_info,
         )
         step = await self._postprocess(step, **kwargs)
 
-        return AgentFlowOutput(steps=[step], metrics=metrics)
+        return AgentFlowOutput(steps=[step], metrics=metrics, **terminal_status(generation_info))
diff --git a/agent_r1/env/envs/tool.py b/agent_r1/env/envs/tool.py
index a13df10..5711bcb 100644
--- a/agent_r1/env/envs/tool.py
+++ b/agent_r1/env/envs/tool.py
@@ -79,7 +79,7 @@ async def step(self, action: Action) -> tuple[Observation, float, bool, dict[str
         self._messages.append({"role": "assistant", "content": action.text})
 
         if not tool_calls:
-            return Observation(messages=list(self._messages)), None, True, {}
+            return Observation(messages=list(self._messages)), None, True, {"termination_reason": "final_answer"}
 
         async def _execute_one(
             tc: ToolCallAction,
diff --git a/docs/tutorials/remax-greedy-rollout.md b/docs/tutorials/remax-greedy-rollout.md
new file mode 100644
index 0000000..f5b5269
--- /dev/null
+++ b/docs/tutorials/remax-greedy-rollout.md
@@ -0,0 +1,91 @@
+# ReMax: complete greedy rollout collection
+
+This first stage provides collection APIs, not a complete ReMax training path.
+The trainer still uses its existing sampler; setting `algorithm.adv_estimator=remax`
+still raises `NotImplementedError`. Advantage computation and actor loss integration
+are a separate next step.
+
+## Using the collection APIs
+
+Use the initialized `AgentFlowManager` with an **unrepeated** generation batch:
+
+```python
+# Inside an initialized trainer, before gen_batch.repeat(...).
+# Use the normal _get_gen_batch preparation and task metadata.
+manager = trainer.async_rollout_manager
+greedy = manager.generate_greedy_sequences(gen_batch)
+
+# Alternatively, collect both paths and pair their rewards in one call.
+collection = manager.collect_remax_rollouts(gen_batch, num_samples=2)
+sampled = collection.sampled
+greedy = collection.greedy
+baseline_per_sampled_step = sampled.batch["reward_baselines"]
+baseline_by_task = collection.baselines
+```
+
+These are two alternative entry points, not two calls needed for one collection.
+`num_samples` defaults to `actor_rollout_ref.rollout.n`. Collection runs one greedy
+episode per original task, then `num_samples` independently reset sampled episodes.
+Do not update actor weights or change environment/reward configuration between these
+passes. No advantage, return, or actor update is computed here.
+
+Each greedy step generates with request-local `temperature=0`, `top_p=1`, and
+`logprobs=False`. Validation temperature cannot override greedy mode. Both sampling
+and greedy requests now explicitly set `max_tokens=rollout.response_length`, so the
+backend budget matches the response slice used by the flow. This replaces the previous
+implicit backend budget; it can affect sampled episodes that formerly generated more
+tokens than their locally retained response. Global sampling settings are unchanged.
+Backend numerical nondeterminism and stochastic external environments are not removed
+by greedy token selection.
+
+## Inspecting a complete trajectory
+
+Outputs remain flattened step-level `DataProto` objects. In `non_tensor_batch`:
+
+- `source_uid` identifies the original task; duplicate original `uid` values are rejected.
+- `trajectory_uids` identifies an episode, and `step_indices` orders its steps.
+- `rollout_mode` distinguishes `greedy` from `sample`.
+- `terminated`, `truncated`, and `termination_reason` are episode end markers on the
+  last step only; previous steps have `False`, `False`, and `ongoing`.
+- `generation_stop_reason`, `response_at_limit`, and `response_truncated` describe
+  each model generation.
+
+Group rows by `trajectory_uids`, then sort by `step_indices` to inspect that episode's
+prompts, actions, and immediate rewards. The next prompt is built from feedback to
+that episode's own action; greedy does not reuse sampled tool feedback.
+
+The baseline is the undiscounted sum of all existing step rewards in that greedy
+episode (`rm_scores` masked by `response_mask`), not just the last step. Pairing uses
+`source_uid`, never row order or equal episode length. The resulting scalar is
+broadcast to every sampled step for that task. It is a baseline, **not an advantage**.
+
+## End states and failures
+
+Built-in single-step, generic environment, HotpotQA, PaperSearch, ALFWorld, and WebShop
+flows report end states. Natural final answers/environment `done` are distinguished
+from `max_steps`, `prompt_length`, response budget, and environment time limits.
+`terminated` means the episode ended naturally, not that the task succeeded.
+Existing recipe reward rules, including ALFWorld's synthetic final reward row, are
+preserved; baseline summation uses those existing rows.
+
+An explicitly truncated episode keeps its existing finite-horizon reward. Some verl
+backends collapse `stop` and `length` into `completed`; without a raw finish reason,
+full-budget responses without a final EOS are conservatively flagged as truncated.
+This fallback can also flag a natural stop at the exact budget.
+
+Empty/aborted generations, missing or duplicate greedy episodes, noncontiguous steps,
+unknown episode end status, and nonfinite rewards raise errors. They are not silently
+replaced with a zero baseline. A custom flow must declare an explicit end status and
+propagate per-generation metadata to participate in collection.
+
+## Verification
+
+```bash
+python3 -B -m unittest discover -s tests -v
+```
+
+Dependency-light tests execute the production orchestration and environment-loop
+method bodies with explicit test doubles. Optional real DataProto/Pydantic/CPU-tensor
+tests run when the verl/PyTorch stack is installed. Neither suite is a GPU model or
+live recipe environment integration test. Verify on a configured training machine
+before connecting collection to advantage computation and actor updates.
diff --git a/docs/zh/tutorials/remax-greedy-rollout.md b/docs/zh/tutorials/remax-greedy-rollout.md
new file mode 100644
index 0000000..46fa397
--- /dev/null
+++ b/docs/zh/tutorials/remax-greedy-rollout.md
@@ -0,0 +1,84 @@
+# ReMax：完整贪心轨迹采集
+
+当前完成的是第一阶段的**采集接口**，不是完整 ReMax 训练链路。
+现有 trainer 仍使用原来的采样入口；设置 `algorithm.adv_estimator=remax`
+仍会抛出 `NotImplementedError`。advantage 计算和 actor loss 接入留到下一阶段。
+
+## 使用入口
+
+在已初始化的 trainer 中，把**尚未按 rollout.n repeat 的原始任务生成批次**传入：
+
+```python
+# gen_batch 来自现有 _get_gen_batch 流程，保留正常的任务元数据。
+manager = trainer.async_rollout_manager
+
+# 入口一：只获取每个原始任务的一条完整 greedy 轨迹。
+greedy = manager.generate_greedy_sequences(gen_batch)
+
+# 入口二：一次完成 greedy + sampled 采集和 baseline 配对。
+collection = manager.collect_remax_rollouts(gen_batch, num_samples=2)
+sampled = collection.sampled
+greedy = collection.greedy
+baseline_per_sampled_step = sampled.batch["reward_baselines"]
+baseline_by_task = collection.baselines
+```
+
+两个入口任选其一，不需要为一次采集同时调用。`num_samples` 不指定时采用
+`actor_rollout_ref.rollout.n`。组合入口先对 B 个任务各采集一条 greedy episode，
+再采集 B × num_samples 条 sampled episode；每条都重新创建 flow、独立 reset 环境。
+两次采集之间不能更新 actor 权重或更改环境、奖励配置。接口不计算 advantage/return，
+也不触发 actor 更新。
+
+greedy 的每一步请求使用 `temperature=0`、`top_p=1`、`logprobs=False`，不会更改
+全局采样配置，验证集温度也不会覆盖 greedy 模式。greedy 和 sampled 都显式传入
+`max_tokens=rollout.response_length`，确保后端生成预算和 flow 保留的响应长度一致。
+这替代了此前的后端隐式预算，可能影响原来先超长生成、再在本地截取的 sampled 轨迹。
+贪心 token 选择不保证消除后端数值不确定性或外部环境随机性。
+
+## 查看完整轨迹与 baseline
+
+输出仍是按 step 展平的 `DataProto`。`non_tensor_batch` 新增或保留：
+
+- `source_uid`：原始任务标识，用于 sampled/greedy 配对，原始任务不允许重复 `uid`。
+- `trajectory_uids`、`step_indices`：episode 标识及步序号。
+- `rollout_mode`：`greedy` 或 `sample`。
+- `terminated`、`truncated`、`termination_reason`：只在 episode 最后一行标记结束状态；
+  前面的行统一为 `False`、`False`、`ongoing`。
+- `generation_stop_reason`、`response_at_limit`、`response_truncated`：每一步的生成状态。
+
+按 `trajectory_uids` 分组，再按 `step_indices` 排序，可还原该 episode 各步的 prompt、
+action 和即时奖励。下一步 prompt 来自当前 episode 自己执行 action 后获得的反馈，
+不会复用 sampled 的工具反馈。
+
+baseline 使用 greedy **所有 step 的原有即时奖励之和**，即对 `rm_scores` 按
+`response_mask` 掩码求和后，再对整条 episode 求和，当前采用不折扣的累计奖励。
+它不是仅取最后一步，也不要求 sampled 和 greedy 步数相同。
+通过 `source_uid` 将一个任务的 baseline 广播到它的每个 sampled step，
+结果位于 `sampled.batch["reward_baselines"]`，此时仍是 baseline，**不是 advantage**。
+
+## 结束状态与异常处理
+
+单步 flow、通用环境循环、HotpotQA、PaperSearch、ALFWorld、WebShop 已补充结束状态。
+自然 final answer/环境 done 与 `max_steps`、`prompt_length`、响应预算、环境时间限制
+截断分开记录。`terminated` 表示 episode 自然结束，不代表任务成功。
+保持各 recipe 的现有奖励规则，包括 ALFWorld 原有的合成最终奖励行；baseline 按这些
+现有 step 行累计。
+
+明确截断的 episode 继续使用现有有限步数下获得的奖励。部分 verl 后端把 stop/length
+都映射成 completed：缺少原始 finish reason 时，恰好耗尽预算且末尾不是 EOS 的响应
+会保守地标记为截断；恰好在预算边界自然停止的响应也可能被这样标记。
+
+空响应、生成中止、缺失/重复 greedy episode、步序不连续、未知结束状态、非有限奖励
+都会报错，不会偷偷用 0 baseline 代替。自定义 flow 需要显式设置结束状态，并记录
+逐步生成元数据后才能参与这一采集流程。
+
+## 验证范围
+
+```bash
+python3 -B -m unittest discover -s tests -v
+```
+
+轻量测试使用测试替身，执行生产代码的采集编排与环境循环方法。
+安装 verl/PyTorch 依赖后还可执行真实 DataProto/Pydantic/CPU tensor 测试。
+这些都不是 GPU 模型或真实任务环境的端到端验证；接入训练前需要在配置完整的训练机
+上确认实际生成、环境隔离及奖励行为。
diff --git a/mkdocs.yml b/mkdocs.yml
index e0e08a2..97601d0 100644
--- a/mkdocs.yml
+++ b/mkdocs.yml
@@ -49,6 +49,7 @@ nav:
       - tutorials/index.md
       - Agent Task Tutorial: tutorials/agent-task.md
       - Recipes and Algorithms: tutorials/recipes-and-algorithms.md
+      - ReMax Greedy Rollout Collection: tutorials/remax-greedy-rollout.md
   - 中文:
       - zh/index.md
       - 快速开始:
@@ -63,6 +64,7 @@ nav:
           - zh/tutorials/index.md
           - 智能体任务教程: zh/tutorials/agent-task.md
           - Recipes 与算法: zh/tutorials/recipes-and-algorithms.md
+          - ReMax 贪心轨迹采集: zh/tutorials/remax-greedy-rollout.md
 
 plugins:
   - search
diff --git a/recipes/alfworld/alfworld_agent_flow.py b/recipes/alfworld/alfworld_agent_flow.py
index 30bc030..8ee4f5c 100644
--- a/recipes/alfworld/alfworld_agent_flow.py
+++ b/recipes/alfworld/alfworld_agent_flow.py
@@ -123,6 +123,7 @@ async def run(self, sampling_params: dict[str, Any], **kwargs) -> AgentFlowOutpu
         done = False
         final_success_flag: bool | None = None
         dense_reward_sum = 0.0
+        end_status = {"terminated": False, "truncated": True, "termination_reason": "max_steps"}
 
         def build_reward_extra_info(step_env_reward: float = 0.0) -> dict[str, Any]:
             return {
@@ -159,6 +160,7 @@ def build_reward_extra_info(step_env_reward: float = 0.0) -> dict[str, Any]:
                 )
 
             response_ids = output.token_ids[: self.response_length]
+            generation_info = self._generation_metadata(output)
             _, tool_calls = await self.tool_parser.extract_tool_calls(response_ids)
             if not tool_calls:
                 response_text = self.tokenizer.decode(response_ids, skip_special_tokens=True)
@@ -193,6 +195,13 @@ def build_reward_extra_info(step_env_reward: float = 0.0) -> dict[str, Any]:
                         env_reward = float(result["reward"])
                         done = bool(result["done"])
                         info = result.get("info", {}) or {}
+                        if done:
+                            env_truncated = bool(info.get("truncated", info.get("TimeLimit.truncated", False)))
+                            end_status = {
+                                "terminated": not env_truncated,
+                                "truncated": env_truncated,
+                                "termination_reason": "env_truncated" if env_truncated else "env_done",
+                            }
                         admissible_commands = info.get("admissible_commands")
                         self.current_admissible_commands = (
                             admissible_commands if isinstance(admissible_commands, list) else []
@@ -224,6 +233,7 @@ def build_reward_extra_info(step_env_reward: float = 0.0) -> dict[str, Any]:
                     },
                 },
             )
+            step.extra_fields.update(generation_info)
             step = await self._postprocess(step, **kwargs)
             self.steps.append(step)
 
@@ -238,8 +248,9 @@ def build_reward_extra_info(step_env_reward: float = 0.0) -> dict[str, Any]:
                         "reward_extra_info": build_reward_extra_info(),
                     },
                 )
+                final_step.extra_fields.update(generation_info)
                 final_step = await self._postprocess(final_step, **kwargs)
                 self.steps.append(final_step)
                 break
 
-        return AgentFlowOutput(steps=self.steps, metrics=metrics)
+        return AgentFlowOutput(steps=self.steps, metrics=metrics, **end_status)
diff --git a/recipes/alfworld/env/alfworld_wrapper.py b/recipes/alfworld/env/alfworld_wrapper.py
index 72f02ea..8fa1eb9 100644
--- a/recipes/alfworld/env/alfworld_wrapper.py
+++ b/recipes/alfworld/env/alfworld_wrapper.py
@@ -137,7 +137,10 @@ def _normalize_step_output(self, step_output: Any) -> tuple[Any, float, bool, di
             terminated = self._unwrap_batch_item(terminated)
             truncated = self._unwrap_batch_item(truncated)
             info = self._unwrap_batch_item(info)
-            return obs, float(reward), bool(terminated or truncated), dict(info or {})
+            info = dict(info or {})
+            # Preserve Gymnasium's distinction through the legacy done interface.
+            info.update(terminated=bool(terminated), truncated=bool(truncated))
+            return obs, float(reward), bool(terminated or truncated), info
         if isinstance(step_output, tuple) and len(step_output) == 4:
             obs, reward, done, info = step_output
             reward = self._unwrap_batch_item(reward)
diff --git a/recipes/hotpotqa/hotpotqa_agent_flow.py b/recipes/hotpotqa/hotpotqa_agent_flow.py
index 7529893..de196a2 100755
--- a/recipes/hotpotqa/hotpotqa_agent_flow.py
+++ b/recipes/hotpotqa/hotpotqa_agent_flow.py
@@ -23,6 +23,7 @@
 from transformers import AutoProcessor, AutoTokenizer
 
 from agent_r1.agent_flow.agent_flow import AgentFlowBase, AgentFlowOutput, AgentFlowStep, register
+from agent_r1.agent_flow.rollout_utils import terminal_status
 from agent_r1.reward_loop.reward_loop import RewardLoopWorker
 from recipes.hotpotqa.env.search_tool import (
     DEFAULT_HOTPOTQA_EMBEDDING_MODEL,
@@ -305,6 +306,7 @@ async def run(self, sampling_params: dict[str, Any], **kwargs) -> AgentFlowOutpu
             self._do_search(question, passages, history_actions)
 
         tool_feedback_lines: list[str] = []
+        end_status = {"terminated": False, "truncated": True, "termination_reason": "max_steps"}
 
         num_steps = 0
         while num_steps < self.max_steps:
@@ -326,6 +328,7 @@ async def run(self, sampling_params: dict[str, Any], **kwargs) -> AgentFlowOutpu
                 )
 
             response_ids = output.token_ids[: self.response_length]
+            generation_info = self._generation_metadata(output)
             _, tool_calls = await self.tool_parser.extract_tool_calls(response_ids)
             response_text = self.tokenizer.decode(response_ids, skip_special_tokens=True)
 
@@ -352,6 +355,7 @@ async def run(self, sampling_params: dict[str, Any], **kwargs) -> AgentFlowOutpu
                     reward_score=None,
                     extra_fields=self._make_extra_fields(anchor_obs, history_actions),
                 )
+                step.extra_fields.update(generation_info)
                 step = await self._postprocess(step, **kwargs)
                 ri = step.extra_fields.get("reward_extra_info", {})
                 step.extra_fields["reward_extra_info"] = {
@@ -359,6 +363,7 @@ async def run(self, sampling_params: dict[str, Any], **kwargs) -> AgentFlowOutpu
                     "acc": ri.get("acc", 0.0),
                 }
                 steps.append(step)
+                end_status = terminal_status(generation_info)
                 break
 
             tool_calls = tool_calls[: self.max_parallel_calls]
@@ -390,10 +395,11 @@ async def run(self, sampling_params: dict[str, Any], **kwargs) -> AgentFlowOutpu
                 reward_score=0.0,
                 extra_fields=self._make_extra_fields(anchor_obs, history_actions),
             )
+            step.extra_fields.update(generation_info)
             step = await self._postprocess(step, **kwargs)
             steps.append(step)
 
-        return AgentFlowOutput(steps=steps, metrics=metrics)
+        return AgentFlowOutput(steps=steps, metrics=metrics, **end_status)
 
     def _do_search(self, query: str, passages: list[tuple[str, str]], history_actions: list[str]) -> None:
         """Execute a single search query and update state."""
diff --git a/recipes/paper_search/paper_search_agent_flow.py b/recipes/paper_search/paper_search_agent_flow.py
index 5d25868..4de69fe 100755
--- a/recipes/paper_search/paper_search_agent_flow.py
+++ b/recipes/paper_search/paper_search_agent_flow.py
@@ -6,6 +6,7 @@
 from transformers import AutoProcessor, AutoTokenizer
 
 from agent_r1.agent_flow.agent_flow import AgentFlowBase, AgentFlowOutput, AgentFlowStep, register
+from agent_r1.agent_flow.rollout_utils import terminal_status
 from agent_r1.reward_loop.reward_loop import RewardLoopWorker
 from recipes.paper_search.env.paper_client import PaperSearchClient, SelectorClient
 from recipes.paper_search.prompts import PAPERSEARCH_TOOL_SCHEMAS
@@ -83,6 +84,7 @@ async def run(self, sampling_params: dict[str, Any], **kwargs) -> AgentFlowOutpu
         total_search_action_count = 0
         total_expand_action_count = 0
         num_steps = 0
+        end_status = {"terminated": False, "truncated": True, "termination_reason": "max_steps"}
 
         while num_steps < self.max_steps:
             num_steps += 1
@@ -98,6 +100,7 @@ async def run(self, sampling_params: dict[str, Any], **kwargs) -> AgentFlowOutpu
                 )
 
             response_ids = output.token_ids[: self.response_length]
+            generation_info = self._generation_metadata(output)
             _, tool_calls = await self.tool_parser.extract_tool_calls(response_ids)
             response_text = self.tokenizer.decode(response_ids, skip_special_tokens=True)
 
@@ -118,8 +121,10 @@ async def run(self, sampling_params: dict[str, Any], **kwargs) -> AgentFlowOutpu
                         },
                     },
                 )
+                step.extra_fields.update(generation_info)
                 step = await self._postprocess(step, **kwargs)
                 self.steps.append(step)
+                end_status = terminal_status(generation_info, reason="no_tool_calls")
                 break
 
             with simple_timer("tool_calls", metrics):
@@ -140,7 +145,8 @@ async def run(self, sampling_params: dict[str, Any], **kwargs) -> AgentFlowOutpu
                     },
                 },
             )
+            step.extra_fields.update(generation_info)
             step = await self._postprocess(step, **kwargs)
             self.steps.append(step)
 
-        return AgentFlowOutput(steps=self.steps, metrics=metrics)
+        return AgentFlowOutput(steps=self.steps, metrics=metrics, **end_status)
diff --git a/recipes/webshop/webshop_agent_flow.py b/recipes/webshop/webshop_agent_flow.py
index d3b369d..935acb3 100644
--- a/recipes/webshop/webshop_agent_flow.py
+++ b/recipes/webshop/webshop_agent_flow.py
@@ -109,6 +109,7 @@ async def run(self, sampling_params: dict[str, Any], **kwargs) -> AgentFlowOutpu
         final_task_score = 0.0
         final_info: dict[str, Any] = {}
         num_steps = 0
+        end_status = {"terminated": False, "truncated": True, "termination_reason": "max_steps"}
 
         while num_steps < self.max_steps and not done:
             num_steps += 1
@@ -129,6 +130,7 @@ async def run(self, sampling_params: dict[str, Any], **kwargs) -> AgentFlowOutpu
                 )
 
             response_ids = output.token_ids[: self.response_length]
+            generation_info = self._generation_metadata(output)
             _, tool_calls = await self.tool_parser.extract_tool_calls(response_ids)
             if not tool_calls:
                 response_text = self.tokenizer.decode(response_ids, skip_special_tokens=True)
@@ -166,6 +168,13 @@ async def run(self, sampling_params: dict[str, Any], **kwargs) -> AgentFlowOutpu
                     env_reward = float(result["reward"])
                     done = bool(result["done"])
                     step_info = result.get("info") or {}
+                    if done:
+                        env_truncated = bool(step_info.get("truncated", step_info.get("TimeLimit.truncated", False)))
+                        end_status = {
+                            "terminated": not env_truncated,
+                            "truncated": env_truncated,
+                            "termination_reason": "env_truncated" if env_truncated else "env_done",
+                        }
                     success = bool(step_info.get("success", env_reward >= 0.999))
                     step_reward = self.success_reward if done and success else 0.0
                     available_actions = step_info.get("available_actions", available_actions)
@@ -212,10 +221,11 @@ async def run(self, sampling_params: dict[str, Any], **kwargs) -> AgentFlowOutpu
                     "reward_extra_info": reward_extra_info,
                 },
             )
+            step.extra_fields.update(generation_info)
             step = await self._postprocess(step, **kwargs)
             self.steps.append(step)
 
             if done:
                 break
 
-        return AgentFlowOutput(steps=self.steps, metrics=metrics)
+        return AgentFlowOutput(steps=self.steps, metrics=metrics, **end_status)
diff --git a/tests/test_greedy_rollouts.py b/tests/test_greedy_rollouts.py
new file mode 100644
index 0000000..642f46c
--- /dev/null
+++ b/tests/test_greedy_rollouts.py
@@ -0,0 +1,593 @@
+"""CPU-only contract tests using simulated generation/environment backends.
+
+Run: python3 -B -m unittest discover -s tests -v
+
+The production module imports the GPU/Ray stack. To exercise its actual loop
+and orchestration methods on a dependency-free machine, load those methods
+from their AST with explicit test doubles. This is NOT a GPU integration test.
+"""
+
+import ast
+import asyncio
+import importlib.util
+import math
+import sys
+import unittest
+from contextlib import nullcontext
+from copy import deepcopy
+from itertools import zip_longest
+from pathlib import Path
+from types import SimpleNamespace
+from uuid import uuid4
+
+ROOT = Path(__file__).resolve().parents[1]
+spec = importlib.util.spec_from_file_location(
+    "greedy_rollout_utils_under_test", ROOT / "agent_r1/agent_flow/rollout_utils.py"
+)
+utils = importlib.util.module_from_spec(spec)
+sys.modules[spec.name] = utils
+spec.loader.exec_module(utils)
+
+
+def load_methods(path, class_name, method_names, namespace):
+    """Compile unmodified production method bodies, removing only decorators."""
+    tree = ast.parse((ROOT / path).read_text())
+    definition = next(node for node in tree.body if isinstance(node, ast.ClassDef) and node.name == class_name)
+    nodes = []
+    for method in definition.body:
+        if isinstance(method, (ast.FunctionDef, ast.AsyncFunctionDef)) and method.name in method_names:
+            method = deepcopy(method)
+            method.decorator_list = []
+            nodes.append(method)
+    if len(nodes) != len(method_names):
+        raise AssertionError(f"Missing production methods in {class_name}")
+    module = ast.Module(
+        body=[ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0), *nodes],
+        type_ignores=[],
+    )
+    ast.fix_missing_locations(module)
+    scope = dict(namespace)
+    exec(compile(module, str(ROOT / path), "exec"), scope)
+    return {name: scope[name] for name in method_names}
+
+
+def rollout_config():
+    return SimpleNamespace(
+        temperature=0.8,
+        top_p=0.9,
+        calculate_log_probs=True,
+        response_length=8,
+        n=2,
+        val_kwargs=SimpleNamespace(temperature=0.4, top_p=0.7),
+        agent=SimpleNamespace(default_agent_flow="fake"),
+    )
+
+
+def zip_with_strict(*iterables, strict=False):
+    """Python 3.9 test-host shim for the production Python 3.10+ zip contract."""
+    if not strict:
+        yield from zip(*iterables)  # noqa: B905
+        return
+    sentinel = object()
+    for row in zip_longest(*iterables, fillvalue=sentinel):
+        if any(value is sentinel for value in row):
+            raise ValueError("zip() arguments have different lengths")
+        yield row
+
+
+def baseline_row(uid="a", trajectory="g-a", step=0, reward=1.0, final=True, truncated=False):
+    return {
+        "source_uid": uid,
+        "trajectory_uid": trajectory,
+        "step_index": step,
+        "reward": reward,
+        "rollout_mode": "greedy",
+        "terminated": final and not truncated,
+        "truncated": final and truncated,
+        "termination_reason": ("max_steps" if truncated else "env_done") if final else "ongoing",
+    }
+
+
+class SamplingAndPairingTests(unittest.TestCase):
+    def test_sampling_and_validation_defaults(self):
+        config = rollout_config()
+        params = utils.build_sampling_params(config, {})
+        self.assertEqual((params["temperature"], params["top_p"]), (0.8, 0.9))
+        self.assertEqual(params["max_tokens"], 8)
+        self.assertTrue(params["logprobs"])
+        params = utils.build_sampling_params(config, {"validate": True})
+        self.assertEqual((params["temperature"], params["top_p"]), (0.4, 0.7))
+
+    def test_greedy_overrides_validation_without_mutation(self):
+        config = rollout_config()
+        for validate in (False, True):
+            with self.subTest(validate=validate):
+                meta = {"rollout_mode": "greedy", "validate": validate}
+                params = utils.build_sampling_params(config, meta)
+                self.assertEqual((params["temperature"], params["top_p"]), (0.0, 1.0))
+                self.assertEqual(params["repetition_penalty"], 1.0)
+                self.assertFalse(params["logprobs"])
+                self.assertNotIn("do_sample", params)
+                self.assertNotIn("n", params)
+                params["temperature"] = 99
+        self.assertEqual(config.temperature, 0.8)
+        self.assertEqual(config.val_kwargs.temperature, 0.4)
+
+    def test_bad_modes_and_budgets_fail(self):
+        with self.assertRaises(ValueError):
+            utils.build_sampling_params(rollout_config(), {"rollout_mode": "typo"})
+        config = rollout_config()
+        config.response_length = 0
+        with self.assertRaises(ValueError):
+            utils.build_sampling_params(config, {})
+
+    def test_generation_empty_and_aborted_fail(self):
+        for tokens, reason in (([], "completed"), ([1], "aborted"), ([1], "error")):
+            with self.subTest(reason=reason), self.assertRaises(RuntimeError):
+                utils.generation_metadata(SimpleNamespace(token_ids=tokens, stop_reason=reason), 8)
+
+    def test_budget_and_eos_classification(self):
+        cases = [
+            ([1], None, "completed", False),
+            ([1, 2], None, "completed", True),
+            ([1, 9], None, "completed", False),
+            ([1, 9, 3], None, "stop", True),
+            ([1, 2], "stop", "completed", False),
+            ([1], "length", "completed", True),
+        ]
+        for tokens, finish, stop, truncated in cases:
+            with self.subTest(tokens=tokens, finish=finish):
+                output = SimpleNamespace(token_ids=tokens, finish_reason=finish, stop_reason=stop)
+                meta = utils.generation_metadata(output, 2, eos_token_id=[9, 10])
+                self.assertEqual(meta["response_truncated"], truncated)
+                status = utils.terminal_status(meta)
+                self.assertEqual(status["truncated"], truncated)
+                self.assertEqual(status["terminated"], not truncated)
+
+    def test_multistep_rewards_pair_by_uid_not_order(self):
+        rows = [
+            baseline_row(step=1, reward=0.75),
+            baseline_row(uid="b", trajectory="g-b", reward=2, truncated=True),
+            baseline_row(step=0, reward=-0.25, final=False),
+        ]
+        baselines = utils.summarize_greedy_baselines(rows, ["a", "b"])
+        self.assertEqual(baselines["a"].reward, 0.5)
+        self.assertEqual(baselines["a"].num_steps, 2)
+        self.assertTrue(baselines["b"].truncated)
+        self.assertEqual(utils.pair_baseline_rewards(["b", "a", "b", "a", "a"], baselines), [2, 0.5, 2, 0.5, 0.5])
+
+    def test_missing_duplicate_mixed_and_invalid_baselines_fail(self):
+        bad_rows = [
+            [],
+            [baseline_row(uid="other")],
+            [baseline_row(), baseline_row(trajectory="second")],
+            [baseline_row(step=1)],
+            [baseline_row(), baseline_row()],
+            [baseline_row(final=False)],
+            [baseline_row(), baseline_row(step=1)],
+            [baseline_row(final=False), baseline_row(uid="b", step=1)],
+            [{**baseline_row(), "rollout_mode": "sample"}],
+            [{**baseline_row(), "termination_reason": "unknown"}],
+            [{**baseline_row(), "truncated": True}],
+        ]
+        for rows in bad_rows:
+            with self.subTest(rows=rows), self.assertRaises(ValueError):
+                utils.summarize_greedy_baselines(rows, ["a"])
+
+    def test_nonfinite_rewards_fail(self):
+        for reward in (math.nan, math.inf, -math.inf):
+            with self.subTest(reward=reward), self.assertRaises(ValueError):
+                utils.summarize_greedy_baselines([baseline_row(reward=reward)], ["a"])
+
+    def test_bad_source_uids_and_missing_pair_fail(self):
+        for uid in (None, "", "  "):
+            with self.subTest(uid=uid), self.assertRaises(ValueError):
+                utils.normalize_source_uid(uid)
+        self.assertEqual(utils.normalize_source_uid(42), "42")
+        with self.assertRaises(ValueError):
+            utils.summarize_greedy_baselines([baseline_row()], ["a", "a"])
+        with self.assertRaises(ValueError):
+            utils.pair_baseline_rewards(["a"], {})
+
+
+class FakeArray(list):
+    def tolist(self):
+        return list(self)
+
+
+class FakeTensor:
+    def __init__(self, values, dtype="float32", device="cpu"):
+        self.values = values
+        self.dtype = dtype
+        self.device = device
+
+    def __mul__(self, other):
+        return FakeTensor([[a * b for a, b in zip(x, y)] for x, y in zip(self.values, other.values)])
+
+    def sum(self, dim):
+        assert dim == -1
+        return FakeTensor([sum(row) for row in self.values])
+
+    def tolist(self):
+        return deepcopy(self.values)
+
+
+class FakeDataProto:
+    def __init__(self, uids, payloads=None, meta_info=None, batch=None, non_tensor_batch=None):
+        self.non_tensor_batch = non_tensor_batch or {"uid": FakeArray(uids)}
+        if payloads is not None:
+            self.non_tensor_batch["payload"] = FakeArray(payloads)
+        self.meta_info = meta_info or {}
+        self.batch = batch or {}
+
+    def __len__(self):
+        return len(next(iter(self.non_tensor_batch.values())))
+
+    def repeat(self, repeat_times, interleave):
+        assert interleave
+        fields = {key: FakeArray(deepcopy(item) for item in values for _ in range(repeat_times))
+                  for key, values in self.non_tensor_batch.items()}
+        return FakeDataProto([], meta_info=deepcopy(self.meta_info), non_tensor_batch=fields)
+
+    def split(self, split_size):
+        return [
+            FakeDataProto(
+                [], meta_info=deepcopy(self.meta_info),
+                non_tensor_batch={key: FakeArray(values[i:i + split_size])
+                                  for key, values in self.non_tensor_batch.items()},
+            )
+            for i in range(0, len(self), split_size)
+        ]
+
+    @staticmethod
+    def concat(outputs):
+        fields = {key: FakeArray(value for output in outputs for value in output.non_tensor_batch[key])
+                  for key in outputs[0].non_tensor_batch}
+        return FakeDataProto([], non_tensor_batch=fields)
+
+
+class CollectorTests(unittest.TestCase):
+    def setUp(self):
+        namespace = {
+            "deepcopy": deepcopy,
+            "uuid4": uuid4,
+            "np": SimpleNamespace(array=lambda values, dtype: FakeArray(values)),
+            "torch": SimpleNamespace(tensor=lambda values, **kwargs: FakeTensor(values, **kwargs)),
+            "normalize_source_uid": utils.normalize_source_uid,
+            "pair_baseline_rewards": utils.pair_baseline_rewards,
+            "summarize_greedy_baselines": utils.summarize_greedy_baselines,
+            "ReMaxRolloutCollection": utils.ReMaxRolloutCollection,
+        }
+        methods = load_methods(
+            "agent_r1/agent_flow/agent_flow.py", "AgentFlowManager",
+            ["_prepare_original_tasks", "generate_greedy_sequences", "collect_remax_rollouts"], namespace,
+        )
+        methods["_prepare_original_tasks"] = staticmethod(methods["_prepare_original_tasks"])
+        manager_cls = type("TestManager", (), methods)
+        self.manager = manager_cls()
+        self.manager.config = SimpleNamespace(actor_rollout_ref=SimpleNamespace(rollout=rollout_config()))
+        self.calls = []
+        self.bad_greedy = False
+
+        def generate(prompts):
+            self.calls.append(deepcopy(prompts))
+            mode = prompts.meta_info["rollout_mode"]
+            rows = []
+            for i, uid in enumerate(prompts.non_tensor_batch["uid"]):
+                rewards = [0.2, 0.3, 0.5] if i == 0 and mode == "greedy" else [2.0]
+                for step, reward in enumerate(rewards):
+                    rows.append(baseline_row(uid, f"{mode}-{i}", step, reward, final=step == len(rewards) - 1))
+            rows.reverse()  # Simulate arbitrary row order after distributed batching.
+            if self.bad_greedy and mode == "greedy":
+                rows.pop()
+            fields = {
+                key: FakeArray(row[row_key] for row in rows)
+                for key, row_key in {
+                    "source_uid": "source_uid", "trajectory_uids": "trajectory_uid",
+                    "step_indices": "step_index", "terminated": "terminated",
+                    "truncated": "truncated", "termination_reason": "termination_reason",
+                }.items()
+            }
+            fields["rollout_mode"] = FakeArray([mode] * len(rows))
+            batch = {
+                "rm_scores": FakeTensor([[row["reward"], 999] for row in rows]),
+                "response_mask": FakeTensor([[1, 0] for _ in rows]),
+                "responses": FakeTensor([[11, 0] for _ in rows]),
+            }
+            prompts.meta_info["mutated_by_backend"] = True
+            return FakeDataProto([], batch=batch, non_tensor_batch=fields)
+
+        self.manager.generate_sequences = generate
+
+    def test_collect_one_greedy_and_n_samples_with_masked_rewards(self):
+        original = FakeDataProto(["a", "b"], payloads=[{"state": 0}, {"state": 0}])
+        result = self.manager.collect_remax_rollouts(original, num_samples=3)
+        self.assertEqual([len(call) for call in self.calls], [2, 6])
+        self.assertEqual([call.meta_info["rollout_mode"] for call in self.calls], ["greedy", "sample"])
+        self.assertEqual(result.baselines["a"].reward, 1.0)
+        self.assertEqual(result.baselines["b"].reward, 2.0)
+        self.assertEqual(result.sampled.batch["reward_baselines"].tolist(), [2, 2, 2, 1, 1, 1])
+        self.assertNotIn("advantages", result.sampled.batch)
+        self.assertNotIn("reward_baselines", result.greedy.batch)
+        self.assertEqual(original.meta_info, {})
+        self.assertEqual(original.non_tensor_batch["payload"], [{"state": 0}, {"state": 0}])
+
+    def test_generated_uid_is_shared_across_paths_not_written_to_caller(self):
+        original = FakeDataProto([], non_tensor_batch={"payload": FakeArray(["x", "y"])})
+        result = self.manager.collect_remax_rollouts(original)
+        greedy_uids = self.calls[0].non_tensor_batch["uid"]
+        self.assertEqual(len(set(greedy_uids)), 2)
+        self.assertEqual(self.calls[1].non_tensor_batch["uid"], [greedy_uids[0]] * 2 + [greedy_uids[1]] * 2)
+        self.assertEqual(set(result.baselines), set(greedy_uids))
+        self.assertNotIn("uid", original.non_tensor_batch)
+
+    def test_greedy_does_not_repeat_or_mutate_caller(self):
+        original = FakeDataProto(["a", "b"], meta_info={"rollout_mode": "sample", "global_steps": 4})
+        self.manager.generate_greedy_sequences(original)
+        self.assertEqual(len(self.calls[0]), 2)
+        self.assertEqual(original.meta_info, {"rollout_mode": "sample", "global_steps": 4})
+
+    def test_invalid_original_tasks_fail_before_generation(self):
+        for uids in ([], ["a", "a"], [None], [1, "1"]):
+            with self.subTest(uids=uids), self.assertRaises(ValueError):
+                self.manager.collect_remax_rollouts(FakeDataProto(uids))
+        self.assertEqual(self.calls, [])
+
+    def test_bad_sample_count_and_validation_fail_before_generation(self):
+        for count in (0, -1, True, 1.5):
+            with self.subTest(count=count), self.assertRaises(ValueError):
+                self.manager.collect_remax_rollouts(FakeDataProto(["a"]), count)
+        with self.assertRaises(ValueError):
+            self.manager.collect_remax_rollouts(FakeDataProto(["a"], meta_info={"validate": True}))
+        self.assertEqual(self.calls, [])
+
+    def test_incomplete_greedy_fails_without_sampling(self):
+        self.bad_greedy = True
+        with self.assertRaises(ValueError):
+            self.manager.collect_remax_rollouts(FakeDataProto(["a", "b"]))
+        self.assertEqual(len(self.calls), 1)
+
+
+class FakeStep(SimpleNamespace):
+    def __init__(self, **kwargs):
+        super().__init__(extra_fields=kwargs.pop("extra_fields", {}), **kwargs)
+
+
+class FakeFlowOutput(SimpleNamespace):
+    def __init__(self, **kwargs):
+        defaults = dict(source_uid=None, rollout_mode="sample", terminated=False, truncated=False,
+                        termination_reason="unknown")
+        defaults.update(kwargs)
+        super().__init__(**defaults)
+
+
+class LoopAndWorkerTests(unittest.IsolatedAsyncioTestCase):
+    def setUp(self):
+        namespace = {
+            "uuid4": uuid4, "AgentFlowStep": FakeStep, "AgentFlowOutput": FakeFlowOutput,
+            "Action": lambda **kwargs: SimpleNamespace(**kwargs),
+            "simple_timer": lambda *args: nullcontext(),
+            "terminal_status": utils.terminal_status,
+            "logger": SimpleNamespace(warning=lambda *args: None),
+        }
+        methods = load_methods("agent_r1/agent_flow/agent_env_loop.py", "AgentEnvLoop", ["run"], namespace)
+        self.flow_cls = type("ActualLoopHarness", (), methods)
+        self.instances = []
+
+    def make_flow(self, done_after=3, env_truncated=False, final_answer=False):
+        env = SimpleNamespace(position=0, actions=[], reset_count=0)
+
+        def reset(**kwargs):
+            env.position = 0
+            env.reset_count += 1
+            return [0]
+
+        async def step(action):
+            env.actions.append(action.token_ids[0])
+            env.position += 1
+            info = {"truncated": env_truncated}
+            if final_answer:
+                info["termination_reason"] = "final_answer"
+            return [env.position, action.token_ids[0]], float(env.position), env.position >= done_after, info
+
+        env.reset, env.step = reset, step
+        flow = self.flow_cls()
+        flow.prompt_length, flow.response_length, flow.max_steps = 16, 8, 5
+        flow.skip_special_tokens = True
+        flow.tokenizer = SimpleNamespace(decode=lambda tokens, **kwargs: str(tokens))
+        flow.loop = asyncio.get_running_loop()
+        flow.requests = []
+        flow._create_env = lambda **kwargs: env
+
+        async def obs_to_prompt(obs, **kwargs):
+            return obs
+
+        async def generate(**kwargs):
+            flow.requests.append(deepcopy(kwargs))
+            chosen = 11 if kwargs["sampling_params"]["temperature"] == 0 else 21
+            return SimpleNamespace(token_ids=[chosen], log_probs=None, routed_experts=None, stop_reason="completed")
+
+        async def postprocess(step, **kwargs):
+            return step
+
+        flow._obs_to_prompt = obs_to_prompt
+        flow.server_manager = SimpleNamespace(generate=generate)
+        flow._postprocess = postprocess
+        flow._generation_metadata = lambda output: utils.generation_metadata(output, flow.response_length)
+        flow.env = env
+        self.instances.append(flow)
+        return flow
+
+    async def test_actual_loop_generates_every_step_from_its_own_environment(self):
+        greedy = self.make_flow()
+        sampled = self.make_flow(done_after=2)
+        greedy_result = await greedy.run(utils.build_sampling_params(rollout_config(), {"rollout_mode": "greedy"}))
+        sample_result = await sampled.run(utils.build_sampling_params(rollout_config(), {}))
+        self.assertEqual(greedy.env.actions, [11, 11, 11])
+        self.assertEqual(sampled.env.actions, [21, 21])
+        self.assertEqual([request["prompt_ids"] for request in greedy.requests], [[0], [1, 11], [2, 11]])
+        self.assertEqual([step.reward_score for step in greedy_result.steps], [1, 2, 3])
+        self.assertTrue(greedy_result.terminated)
+        self.assertTrue(sample_result.terminated)
+        self.assertEqual(greedy_result.termination_reason, "env_done")
+        self.assertIsNot(greedy.env, sampled.env)
+
+    async def test_actual_loop_max_steps_and_prompt_overflow(self):
+        flow = self.make_flow(done_after=10)
+        flow.max_steps = 2
+        output = await flow.run(utils.build_sampling_params(rollout_config(), {"rollout_mode": "greedy"}))
+        self.assertEqual(len(output.steps), 2)
+        self.assertTrue(output.truncated)
+        self.assertFalse(output.terminated)
+        self.assertEqual(output.termination_reason, "max_steps")
+        flow = self.make_flow()
+        flow.prompt_length = 0
+        output = await flow.run(utils.build_sampling_params(rollout_config(), {}))
+        self.assertEqual(output.steps, [])
+        self.assertEqual(output.termination_reason, "prompt_length")
+        self.assertEqual(flow.requests, [])
+
+    async def test_actual_loop_preserves_environment_time_limit(self):
+        flow = self.make_flow(done_after=1, env_truncated=True)
+        output = await flow.run(utils.build_sampling_params(rollout_config(), {}))
+        self.assertTrue(output.truncated)
+        self.assertEqual(output.termination_reason, "env_truncated")
+
+    async def test_truncated_final_answer_not_claimed_as_natural_end(self):
+        flow = self.make_flow(done_after=1, final_answer=True)
+        flow.response_length = 1
+        output = await flow.run(utils.build_sampling_params(rollout_config(), {}))
+        self.assertTrue(output.truncated)
+        self.assertEqual(output.termination_reason, "response_length")
+
+    async def test_actual_worker_sets_uid_mode_and_creates_fresh_flows(self):
+        async def trajectory_info(step, index, validate):
+            return [dict(step=step, sample_index=i, rollout_n=0, validate=validate) for i in index]
+
+        def instantiate(**kwargs):
+            return self.make_flow(done_after=2)
+
+        namespace = {
+            "asyncio": asyncio, "uuid4": uuid4,
+            "np": SimpleNamespace(
+                array=lambda values, dtype: FakeArray(values), arange=lambda size: FakeArray(range(size))
+            ),
+            "build_sampling_params": utils.build_sampling_params,
+            "normalize_source_uid": utils.normalize_source_uid,
+            "RolloutTraceConfig": SimpleNamespace(
+                get_instance=lambda: SimpleNamespace(max_samples_per_step_per_worker=None)
+            ),
+            "get_trajectory_info": trajectory_info,
+            "rollout_trace_attr": lambda **kwargs: nullcontext(),
+            "_agent_flow_registry": {"fake": {}},
+            "hydra": SimpleNamespace(utils=SimpleNamespace(instantiate=instantiate)),
+            "DictConfigWrap": lambda **kwargs: SimpleNamespace(**kwargs),
+        }
+        methods = load_methods(
+            "agent_r1/agent_flow/agent_flow.py", "AgentFlowWorkerBase",
+            ["generate_sequences", "_run_agent_flow"], namespace,
+        )
+        worker = type("ActualWorkerHarness", (), methods)()
+        worker.config = SimpleNamespace(actor_rollout_ref=SimpleNamespace(rollout=rollout_config()), data={})
+        for name in ("server_manager", "reward_loop_worker", "tokenizer", "processor", "dataset_cls"):
+            setattr(worker, name, None)
+        worker._postprocess = lambda outputs: outputs
+        batch = FakeDataProto(["a", "a", "b"], meta_info={"rollout_mode": "greedy", "validate": True})
+        outputs = await worker.generate_sequences(batch)
+        self.assertEqual([output.source_uid for output in outputs], ["a", "a", "b"])
+        self.assertEqual([output.rollout_mode for output in outputs], ["greedy"] * 3)
+        self.assertEqual(len({id(flow.env) for flow in self.instances}), 3)
+        for flow in self.instances:
+            self.assertEqual(flow.env.reset_count, 1)
+            self.assertEqual(flow.env.actions, [11, 11])
+            self.assertTrue(all(request["sampling_params"]["temperature"] == 0 for request in flow.requests))
+        self.assertEqual(worker.config.actor_rollout_ref.rollout.temperature, 0.8)
+
+
+class DispatchTests(unittest.TestCase):
+    def make_manager(self, fail=False):
+        namespace = {
+            "ray": SimpleNamespace(get=lambda outputs: outputs),
+            "DataProto": FakeDataProto,
+            "zip": zip if sys.version_info >= (3, 10) else zip_with_strict,
+        }
+        methods = load_methods(
+            "agent_r1/agent_flow/agent_flow.py", "AgentFlowManager", ["generate_sequences"], namespace
+        )
+        manager = type("ActualDispatchHarness", (), methods)()
+        self.events = []
+
+        def generate(chunk):
+            self.events.append("generate")
+            if fail:
+                raise RuntimeError("simulated worker failure")
+            chunk.meta_info["metrics"] = [dict(num_steps=1) for _ in range(len(chunk))]
+            return chunk
+
+        manager.agent_flow_workers = [
+            SimpleNamespace(generate_sequences=SimpleNamespace(remote=generate)) for _ in range(8)
+        ]
+        manager.reward_model_manager = SimpleNamespace(
+            wake_up=lambda: self.events.append("reward_wake"), sleep=lambda: self.events.append("reward_sleep")
+        )
+        manager.wake_up = lambda: self.events.append("wake")
+        manager.sleep = lambda: self.events.append("sleep")
+        manager._performance_metrics = lambda *args: {}
+        return manager
+
+    def test_task_count_smaller_than_worker_count(self):
+        manager = self.make_manager()
+        result = manager.generate_sequences(FakeDataProto(["a", "b"]))
+        self.assertEqual(result.meta_info["num_steps"], [1, 1])
+        self.assertEqual(self.events.count("generate"), 2)
+        self.assertEqual(self.events[-2:], ["sleep", "reward_sleep"])
+
+    def test_backend_failure_still_releases_rollout_and_reward_engines(self):
+        manager = self.make_manager(fail=True)
+        with self.assertRaises(RuntimeError):
+            manager.generate_sequences(FakeDataProto(["a"]))
+        self.assertEqual(self.events[-2:], ["sleep", "reward_sleep"])
+
+    def test_rollout_sleep_failure_still_releases_reward_engine(self):
+        manager = self.make_manager()
+
+        def sleep():
+            self.events.append("sleep")
+            raise RuntimeError("simulated sleep failure")
+
+        manager.sleep = sleep
+        with self.assertRaises(RuntimeError):
+            manager.generate_sequences(FakeDataProto(["a"]))
+        self.assertEqual(self.events[-2:], ["sleep", "reward_sleep"])
+
+    def test_empty_input_fails_without_waking_engines(self):
+        manager = self.make_manager()
+        with self.assertRaises(ValueError):
+            manager.generate_sequences(FakeDataProto([]))
+        self.assertEqual(self.events, [])
+
+
+class EnvironmentAdapterTests(unittest.TestCase):
+    def test_alfworld_preserves_gymnasium_end_flags_and_legacy_done(self):
+        methods = load_methods(
+            "recipes/alfworld/env/alfworld_wrapper.py", "AlfworldTextworldEnv",
+            ["_unwrap_batch_item", "_normalize_step_output"], {},
+        )
+        methods["_unwrap_batch_item"] = staticmethod(methods["_unwrap_batch_item"])
+        adapter = type("ActualAdapterHarness", (), methods)()
+        for terminated, truncated in ((False, False), (True, False), (False, True)):
+            with self.subTest(terminated=terminated, truncated=truncated):
+                info = {"success": False}
+                obs, reward, done, result_info = adapter._normalize_step_output(
+                    ("obs", [0.5], [terminated], [truncated], [info])
+                )
+                self.assertEqual((obs, reward, done), ("obs", 0.5, terminated or truncated))
+                self.assertEqual(result_info["terminated"], terminated)
+                self.assertEqual(result_info["truncated"], truncated)
+                self.assertEqual(info, {"success": False})
+        self.assertEqual(adapter._normalize_step_output(("obs", [1], [True], [{}])), ("obs", 1, True, {}))
+        self.assertEqual(adapter._normalize_step_output(("obs", [1], [True])), ("obs", 1, True, {}))
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/test_greedy_tensor_contract.py b/tests/test_greedy_tensor_contract.py
new file mode 100644
index 0000000..a033ca2
--- /dev/null
+++ b/tests/test_greedy_tensor_contract.py
@@ -0,0 +1,75 @@
+"""Optional real DataProto/Pydantic/CPU-tensor tests; never start Ray or a GPU model."""
+
+import importlib.util
+import unittest
+
+DEPENDENCIES = ("torch", "numpy", "pydantic", "ray", "hydra", "transformers", "verl", "tensordict")
+HAS_STACK = all(importlib.util.find_spec(name) is not None for name in DEPENDENCIES)
+
+
+@unittest.skipUnless(HAS_STACK, "Requires installed verl/PyTorch stack for real tensor contract tests")
+class TensorContractTests(unittest.TestCase):
+    def setUp(self):
+        import torch
+
+        from agent_r1.agent_flow.agent_flow import AgentFlowOutput, AgentFlowWorkerBase, _InternalAgentFlowStep
+
+        self.torch = torch
+        self.output_cls = AgentFlowOutput
+        self.step_cls = _InternalAgentFlowStep
+        self.worker = object.__new__(AgentFlowWorkerBase)
+
+    def make_step(self, reward=0.75, mask=None, extra_fields=None):
+        tensor = self.torch.tensor
+        return self.step_cls(
+            prompt_ids=tensor([[1, 2]]), response_ids=tensor([[3, 77, 4]]),
+            input_ids=tensor([[1, 2, 3, 77, 4]]), attention_mask=tensor([[1, 1, 1, 1, 1]]),
+            position_ids=tensor([[0, 1, 2, 3, 4]]), response_mask=tensor([mask or [1, 0, 1]]),
+            reward_score=reward, extra_fields=extra_fields or {},
+        )
+
+    def test_real_flattening_provenance_end_markers_and_last_action_reward(self):
+        outputs = [
+            self.output_cls(
+                steps=[self.make_step(0), self.make_step(0.75, extra_fields={
+                    "source_uid": "spoof", "trajectory_uids": "spoof", "step_indices": 99,
+                })],
+                metrics={}, source_uid="a", rollout_mode="greedy", terminated=True, termination_reason="env_done",
+            ),
+            self.output_cls(
+                steps=[self.make_step(2)], metrics={}, source_uid="b", rollout_mode="greedy",
+                truncated=True, termination_reason="max_steps",
+            ),
+        ]
+        result = self.worker._postprocess(outputs)
+        fields = result.non_tensor_batch
+        self.assertEqual(fields["source_uid"].tolist(), ["a", "a", "b"])
+        self.assertEqual(fields["step_indices"].tolist(), [0, 1, 0])
+        self.assertEqual(fields["terminated"].tolist(), [False, True, False])
+        self.assertEqual(fields["truncated"].tolist(), [False, False, True])
+        self.assertEqual(fields["termination_reason"].tolist(), ["ongoing", "env_done", "max_steps"])
+        self.assertEqual(result.batch["rm_scores"].tolist(), [[0, 0, 0], [0, 0, 0.75], [0, 0, 2]])
+        self.assertEqual(fields["trajectory_uids"][0], fields["trajectory_uids"][1])
+        self.assertNotEqual(fields["trajectory_uids"][1], fields["trajectory_uids"][2])
+
+    def test_real_flattening_rejects_empty_trajectory_and_response(self):
+        for steps in ([], [self.make_step(mask=[0, 0, 0])]):
+            with self.subTest(steps=steps), self.assertRaises(RuntimeError):
+                self.worker._postprocess([self.output_cls(steps=steps, metrics={}, source_uid="a")])
+
+    def test_real_distributed_concat_keeps_new_step_metadata(self):
+        from verl import DataProto
+
+        outputs = []
+        for uid in ("a", "b"):
+            outputs.append(self.worker._postprocess([
+                self.output_cls(steps=[self.make_step()], metrics={}, source_uid=uid, rollout_mode="greedy",
+                                terminated=True, termination_reason="env_done")
+            ]))
+        combined = DataProto.concat(outputs)
+        self.assertEqual(combined.non_tensor_batch["source_uid"].tolist(), ["a", "b"])
+        self.assertEqual(combined.non_tensor_batch["rollout_mode"].tolist(), ["greedy", "greedy"])
+
+
+if __name__ == "__main__":
+    unittest.main()

From fd64792405525e6475e47a68e775c59c161a24c8 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Mon, 21 Sep 2026 09:09:05 +0000
Subject: [PATCH 2/2] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 tests/test_greedy_rollouts.py        | 59 ++++++++++++++++++----------
 tests/test_greedy_tensor_contract.py | 58 ++++++++++++++++++++-------
 2 files changed, 83 insertions(+), 34 deletions(-)

diff --git a/tests/test_greedy_rollouts.py b/tests/test_greedy_rollouts.py
index 642f46c..e6dcd28 100644
--- a/tests/test_greedy_rollouts.py
+++ b/tests/test_greedy_rollouts.py
@@ -225,24 +225,30 @@ def __len__(self):
 
     def repeat(self, repeat_times, interleave):
         assert interleave
-        fields = {key: FakeArray(deepcopy(item) for item in values for _ in range(repeat_times))
-                  for key, values in self.non_tensor_batch.items()}
+        fields = {
+            key: FakeArray(deepcopy(item) for item in values for _ in range(repeat_times))
+            for key, values in self.non_tensor_batch.items()
+        }
         return FakeDataProto([], meta_info=deepcopy(self.meta_info), non_tensor_batch=fields)
 
     def split(self, split_size):
         return [
             FakeDataProto(
-                [], meta_info=deepcopy(self.meta_info),
-                non_tensor_batch={key: FakeArray(values[i:i + split_size])
-                                  for key, values in self.non_tensor_batch.items()},
+                [],
+                meta_info=deepcopy(self.meta_info),
+                non_tensor_batch={
+                    key: FakeArray(values[i : i + split_size]) for key, values in self.non_tensor_batch.items()
+                },
             )
             for i in range(0, len(self), split_size)
         ]
 
     @staticmethod
     def concat(outputs):
-        fields = {key: FakeArray(value for output in outputs for value in output.non_tensor_batch[key])
-                  for key in outputs[0].non_tensor_batch}
+        fields = {
+            key: FakeArray(value for output in outputs for value in output.non_tensor_batch[key])
+            for key in outputs[0].non_tensor_batch
+        }
         return FakeDataProto([], non_tensor_batch=fields)
 
 
@@ -259,8 +265,10 @@ def setUp(self):
             "ReMaxRolloutCollection": utils.ReMaxRolloutCollection,
         }
         methods = load_methods(
-            "agent_r1/agent_flow/agent_flow.py", "AgentFlowManager",
-            ["_prepare_original_tasks", "generate_greedy_sequences", "collect_remax_rollouts"], namespace,
+            "agent_r1/agent_flow/agent_flow.py",
+            "AgentFlowManager",
+            ["_prepare_original_tasks", "generate_greedy_sequences", "collect_remax_rollouts"],
+            namespace,
         )
         methods["_prepare_original_tasks"] = staticmethod(methods["_prepare_original_tasks"])
         manager_cls = type("TestManager", (), methods)
@@ -283,9 +291,12 @@ def generate(prompts):
             fields = {
                 key: FakeArray(row[row_key] for row in rows)
                 for key, row_key in {
-                    "source_uid": "source_uid", "trajectory_uids": "trajectory_uid",
-                    "step_indices": "step_index", "terminated": "terminated",
-                    "truncated": "truncated", "termination_reason": "termination_reason",
+                    "source_uid": "source_uid",
+                    "trajectory_uids": "trajectory_uid",
+                    "step_indices": "step_index",
+                    "terminated": "terminated",
+                    "truncated": "truncated",
+                    "termination_reason": "termination_reason",
                 }.items()
             }
             fields["rollout_mode"] = FakeArray([mode] * len(rows))
@@ -355,8 +366,9 @@ def __init__(self, **kwargs):
 
 class FakeFlowOutput(SimpleNamespace):
     def __init__(self, **kwargs):
-        defaults = dict(source_uid=None, rollout_mode="sample", terminated=False, truncated=False,
-                        termination_reason="unknown")
+        defaults = dict(
+            source_uid=None, rollout_mode="sample", terminated=False, truncated=False, termination_reason="unknown"
+        )
         defaults.update(kwargs)
         super().__init__(**defaults)
 
@@ -364,7 +376,9 @@ def __init__(self, **kwargs):
 class LoopAndWorkerTests(unittest.IsolatedAsyncioTestCase):
     def setUp(self):
         namespace = {
-            "uuid4": uuid4, "AgentFlowStep": FakeStep, "AgentFlowOutput": FakeFlowOutput,
+            "uuid4": uuid4,
+            "AgentFlowStep": FakeStep,
+            "AgentFlowOutput": FakeFlowOutput,
             "Action": lambda **kwargs: SimpleNamespace(**kwargs),
             "simple_timer": lambda *args: nullcontext(),
             "terminal_status": utils.terminal_status,
@@ -468,7 +482,8 @@ def instantiate(**kwargs):
             return self.make_flow(done_after=2)
 
         namespace = {
-            "asyncio": asyncio, "uuid4": uuid4,
+            "asyncio": asyncio,
+            "uuid4": uuid4,
             "np": SimpleNamespace(
                 array=lambda values, dtype: FakeArray(values), arange=lambda size: FakeArray(range(size))
             ),
@@ -484,8 +499,10 @@ def instantiate(**kwargs):
             "DictConfigWrap": lambda **kwargs: SimpleNamespace(**kwargs),
         }
         methods = load_methods(
-            "agent_r1/agent_flow/agent_flow.py", "AgentFlowWorkerBase",
-            ["generate_sequences", "_run_agent_flow"], namespace,
+            "agent_r1/agent_flow/agent_flow.py",
+            "AgentFlowWorkerBase",
+            ["generate_sequences", "_run_agent_flow"],
+            namespace,
         )
         worker = type("ActualWorkerHarness", (), methods)()
         worker.config = SimpleNamespace(actor_rollout_ref=SimpleNamespace(rollout=rollout_config()), data={})
@@ -570,8 +587,10 @@ def test_empty_input_fails_without_waking_engines(self):
 class EnvironmentAdapterTests(unittest.TestCase):
     def test_alfworld_preserves_gymnasium_end_flags_and_legacy_done(self):
         methods = load_methods(
-            "recipes/alfworld/env/alfworld_wrapper.py", "AlfworldTextworldEnv",
-            ["_unwrap_batch_item", "_normalize_step_output"], {},
+            "recipes/alfworld/env/alfworld_wrapper.py",
+            "AlfworldTextworldEnv",
+            ["_unwrap_batch_item", "_normalize_step_output"],
+            {},
         )
         methods["_unwrap_batch_item"] = staticmethod(methods["_unwrap_batch_item"])
         adapter = type("ActualAdapterHarness", (), methods)()
diff --git a/tests/test_greedy_tensor_contract.py b/tests/test_greedy_tensor_contract.py
index a033ca2..9ad4ff3 100644
--- a/tests/test_greedy_tensor_contract.py
+++ b/tests/test_greedy_tensor_contract.py
@@ -22,23 +22,43 @@ def setUp(self):
     def make_step(self, reward=0.75, mask=None, extra_fields=None):
         tensor = self.torch.tensor
         return self.step_cls(
-            prompt_ids=tensor([[1, 2]]), response_ids=tensor([[3, 77, 4]]),
-            input_ids=tensor([[1, 2, 3, 77, 4]]), attention_mask=tensor([[1, 1, 1, 1, 1]]),
-            position_ids=tensor([[0, 1, 2, 3, 4]]), response_mask=tensor([mask or [1, 0, 1]]),
-            reward_score=reward, extra_fields=extra_fields or {},
+            prompt_ids=tensor([[1, 2]]),
+            response_ids=tensor([[3, 77, 4]]),
+            input_ids=tensor([[1, 2, 3, 77, 4]]),
+            attention_mask=tensor([[1, 1, 1, 1, 1]]),
+            position_ids=tensor([[0, 1, 2, 3, 4]]),
+            response_mask=tensor([mask or [1, 0, 1]]),
+            reward_score=reward,
+            extra_fields=extra_fields or {},
         )
 
     def test_real_flattening_provenance_end_markers_and_last_action_reward(self):
         outputs = [
             self.output_cls(
-                steps=[self.make_step(0), self.make_step(0.75, extra_fields={
-                    "source_uid": "spoof", "trajectory_uids": "spoof", "step_indices": 99,
-                })],
-                metrics={}, source_uid="a", rollout_mode="greedy", terminated=True, termination_reason="env_done",
+                steps=[
+                    self.make_step(0),
+                    self.make_step(
+                        0.75,
+                        extra_fields={
+                            "source_uid": "spoof",
+                            "trajectory_uids": "spoof",
+                            "step_indices": 99,
+                        },
+                    ),
+                ],
+                metrics={},
+                source_uid="a",
+                rollout_mode="greedy",
+                terminated=True,
+                termination_reason="env_done",
             ),
             self.output_cls(
-                steps=[self.make_step(2)], metrics={}, source_uid="b", rollout_mode="greedy",
-                truncated=True, termination_reason="max_steps",
+                steps=[self.make_step(2)],
+                metrics={},
+                source_uid="b",
+                rollout_mode="greedy",
+                truncated=True,
+                termination_reason="max_steps",
             ),
         ]
         result = self.worker._postprocess(outputs)
@@ -62,10 +82,20 @@ def test_real_distributed_concat_keeps_new_step_metadata(self):
 
         outputs = []
         for uid in ("a", "b"):
-            outputs.append(self.worker._postprocess([
-                self.output_cls(steps=[self.make_step()], metrics={}, source_uid=uid, rollout_mode="greedy",
-                                terminated=True, termination_reason="env_done")
-            ]))
+            outputs.append(
+                self.worker._postprocess(
+                    [
+                        self.output_cls(
+                            steps=[self.make_step()],
+                            metrics={},
+                            source_uid=uid,
+                            rollout_mode="greedy",
+                            terminated=True,
+                            termination_reason="env_done",
+                        )
+                    ]
+                )
+            )
         combined = DataProto.concat(outputs)
         self.assertEqual(combined.non_tensor_batch["source_uid"].tolist(), ["a", "b"])
         self.assertEqual(combined.non_tensor_batch["rollout_mode"].tolist(), ["greedy", "greedy"])