fix: make rollout sampling deterministic

areal/api/cli_args.py

@@ -222,6 +222,14 @@ class GenerationHyperparameters:
             )
         },
     )
+    seed: int | None = field(
+        default=None,
+        metadata={
+            "help": "Per-request sampling seed sent to the inference backend. Leave "
+            "unset for grouped deterministic rollouts so each sample receives a "
+            "stable, distinct derived seed."
+        },
+    )
     lora_name: str = field(
         default="default_lora",
         metadata={"help": "Lora name to be used for this generation."},
@@ -2107,6 +2115,11 @@ def build_cmd(
         return vLLMConfig.build_cmd_from_args(args)
 
 
+# Keep this list aligned with SGLang's deterministic inference documentation:
+# https://docs.sglang.ai/advanced_features/deterministic_inference.html
+_SGLANG_DETERMINISTIC_ATTENTION_BACKENDS = frozenset({"flashinfer", "fa3", "triton"})
+
+
 @dataclass
 class SGLangConfig:
     """Configuration for SGLang runtime. Refer to:
@@ -2140,6 +2153,7 @@ class SGLangConfig:
     enable_memory_saver: bool = False
     allow_auto_truncate: bool = False
     attention_backend: str | None = "fa3"
+    enable_deterministic_inference: bool = False
     enable_multimodal: bool = False
     sampling_backend: str | None = None
     context_length: int | None = 32768
@@ -2223,6 +2237,19 @@ def build_args(
         node_rank: int = 0,
         pp_size: int = 1,
     ):
+        attention_backend = sglang_config.attention_backend
+        if (
+            sglang_config.enable_deterministic_inference
+            and attention_backend is not None
+            and attention_backend.lower()
+            not in _SGLANG_DETERMINISTIC_ATTENTION_BACKENDS
+        ):
+            logger.warning(
+                "SGLang deterministic inference is only documented for attention "
+                "backends %s; configured attention_backend=%r may be non-deterministic.",
+                sorted(_SGLANG_DETERMINISTIC_ATTENTION_BACKENDS),
+                attention_backend,
+            )
         # Map "all-linear" to "all"
         args: dict = conf_as_dict(sglang_config)
         if sglang_config.enable_multithread_load:
@@ -2457,6 +2484,25 @@ class InferenceEngineConfig:
             "help": "Whether to output verbose tracing messages for each generation request."
         },
     )
+    deterministic_sampling: bool = field(
+        default=False,
+        metadata={
+            "help": "Use stable request seeds for internal OpenAI-proxy/data-proxy "
+            "sessions, canonical group ordering, and task-ID ordering of completed "
+            "rollout results. Concurrent SGLang generation also requires "
+            "sglang.enable_deterministic_inference. End-to-end determinism is only "
+            "supported with max_head_offpolicyness=0."
+        },
+    )
+    serialize_group_samples: bool = field(
+        default=False,
+        metadata={
+            "help": "Run RolloutControllerV2 samples within each group sequentially "
+            "instead of concurrently. This provides stable within-group member "
+            "submission order at the cost of rollout throughput; it does not "
+            "serialize requests across groups."
+        },
+    )
     check_trajectory_format: bool = field(
         default=False,
         metadata={
@@ -2601,6 +2647,13 @@ def __post_init__(self):
             )
         if not self.admin_api_key or not self.admin_api_key.strip():
             raise ValueError("admin_api_key must not be empty or whitespace-only")
+        if self.deterministic_sampling and self.max_head_offpolicyness > 0:
+            logger.warning(
+                "deterministic_sampling=True with max_head_offpolicyness=%d does "
+                "not guarantee deterministic task-to-weight-version mapping; "
+                "set max_head_offpolicyness=0 for end-to-end determinism.",
+                self.max_head_offpolicyness,
+            )
         if (
             self._version == "v2"
             and self.agent is not None
@@ -3352,6 +3405,21 @@ def __post_init__(self):
         """Validate the eval generation config."""
         if self.eval_gconfig is None:
             self.eval_gconfig = self.gconfig.new()
+        if self.rollout.deterministic_sampling:
+            for config_name, generation_config in (
+                ("gconfig", self.gconfig),
+                ("eval_gconfig", self.eval_gconfig),
+            ):
+                if (
+                    generation_config.n_samples > 1
+                    and generation_config.seed is not None
+                ):
+                    raise ValueError(
+                        "deterministic_sampling with grouped rollouts cannot use "
+                        f"a shared {config_name}.seed, because every sample would "
+                        "receive the same sampling seed. Set the seed to null to "
+                        "derive stable per-sample seeds, or set n_samples=1."
+                    )
         if self.gconfig.reward_normalization and self.actor.reward_norm is not None:
             raise ValueError(
                 "gconfig.reward_normalization (rollout-time, per-prompt) and "

areal/engine/sglang_remote.py

@@ -76,6 +76,8 @@ def build_generation_request(
         }
         if stop:
             sample_params["stop"] = stop
+        if gconfig.seed is not None:
+            sample_params["sampling_seed"] = gconfig.seed
 
         payload = {
             "input_ids": req.input_ids.copy(),

areal/engine/vllm_remote.py

@@ -79,6 +79,8 @@ def build_generation_request(
         }
         if gconfig.stop:
             payload["stop"] = gconfig.stop
+        if gconfig.seed is not None:
+            payload["seed"] = gconfig.seed
 
         if with_lora:
             lora_name = gconfig.lora_name

areal/experimental/openai/client.py

@@ -646,6 +646,7 @@ async def create(
         max_total_tokens: int | None | NotGiven = NOT_GIVEN,
         metadata: Metadata | None | NotGiven = NOT_GIVEN,
         n: int | None | NotGiven = NOT_GIVEN,
+        seed: int | None | NotGiven = NOT_GIVEN,
         stop: str | None | list[str] | None | NotGiven = NOT_GIVEN,
         store: bool | None | NotGiven = NOT_GIVEN,
         temperature: float | None | NotGiven = NOT_GIVEN,
@@ -669,6 +670,7 @@ async def create(
         max_total_tokens: int | None | NotGiven = NOT_GIVEN,
         metadata: Metadata | None | NotGiven = NOT_GIVEN,
         n: int | None | NotGiven = NOT_GIVEN,
+        seed: int | None | NotGiven = NOT_GIVEN,
         stop: str | None | list[str] | None | NotGiven = NOT_GIVEN,
         store: bool | None | NotGiven = NOT_GIVEN,
         temperature: float | None | NotGiven = NOT_GIVEN,
@@ -691,6 +693,7 @@ async def create(
         max_total_tokens: int | None | NotGiven = NOT_GIVEN,
         metadata: Metadata | None | NotGiven = NOT_GIVEN,
         n: int | None | NotGiven = NOT_GIVEN,
+        seed: int | None | NotGiven = NOT_GIVEN,
         stop: str | None | list[str] | None | NotGiven = NOT_GIVEN,
         store: bool | None | NotGiven = NOT_GIVEN,
         temperature: float | None | NotGiven = NOT_GIVEN,
@@ -883,6 +886,7 @@ async def create(
             greedy=temp == 0,
             frequency_penalty=frequency_penalty,
             lora_name=self.lora_name,
+            seed=None if is_omitted(seed) else seed,
             stop_token_ids=list(
                 set([self.tokenizer.eos_token_id, self.tokenizer.pad_token_id])
             ),
@@ -1136,6 +1140,7 @@ async def create(
         instructions: str | None | NotGiven = NOT_GIVEN,
         max_output_tokens: int | None | NotGiven = NOT_GIVEN,
         metadata: Metadata | None | NotGiven = NOT_GIVEN,
+        seed: int | None | NotGiven = NOT_GIVEN,
         tool_choice: response_create_params.ToolChoice | NotGiven = NOT_GIVEN,
         tools: Iterable[ToolParam] | NotGiven = NOT_GIVEN,
         temperature: float | None | NotGiven = NOT_GIVEN,
@@ -1291,6 +1296,7 @@ async def create(
             greedy=temp == 0,
             frequency_penalty=frequency_penalty,
             lora_name=self.lora_name,
+            seed=None if is_omitted(seed) else seed,
             stop_token_ids=list(
                 set([self.tokenizer.eos_token_id, self.tokenizer.pad_token_id])
             ),

areal/experimental/openai/proxy/proxy_rollout_server.py

@@ -75,6 +75,10 @@
 _warn_lock = threading.Lock()
 
 
+def _deterministic_sampling_seed(session_id: str, request_index: int) -> int:
+    return seeding.derive_deterministic_seed(session_id, request_index)
+
+
 def _warn_once(msg: str) -> None:
     """Log a warning message, optionally only once if AREAL_PROXY_WARN_ONCE=1."""
     if not _warn_once_enabled:
@@ -126,6 +130,9 @@ def _warn_once(msg: str) -> None:
 _allocated_ports: set[int] = set()
 _port_alloc_lock = asyncio.Lock()
 
+# Deterministic sampling (set from InferenceEngineConfig at setup time).
+_deterministic_sampling: bool = False
+
 # Server config (needed for name_resolve registration)
 _experiment_name: str | None = None
 _trial_name: str | None = None
@@ -270,8 +277,9 @@ async def alloc_ports(raw_request: Request):
 
 def _setup_openai_client():
     global _openai_client, _session_timeout_seconds, _admin_api_key
-    global _message_preprocessors, _prefix_matcher
+    global _message_preprocessors, _prefix_matcher, _deterministic_sampling
     config = _engine.config
+    _deterministic_sampling = bool(getattr(config, "deterministic_sampling", False))
     tokenizer = load_hf_tokenizer(config.tokenizer_path)
     agent_cfg = config.agent
     _openai_client = ArealOpenAI(
@@ -488,6 +496,7 @@ def start_session(request: StartSessionRequest) -> StartSessionResponse:
         _session_cache[session_id] = SessionData(
             session_id=session_id,
             prefix_matcher=_prefix_matcher,
+            sampling_seed_identity=task_id,
         )
         _api_key_to_session[session_api_key] = session_id
         _session_to_api_key[session_id] = session_api_key
@@ -574,8 +583,11 @@ async def _call_client_create(
                 status_code=410, detail=f"Session {session_id} already ended or expired"
             )
         session_data = _session_cache[session_id]
+        session_data.update_last_access()
 
-    session_data.update_last_access()
+    request_index = (
+        session_data.next_sampling_request_index() if _deterministic_sampling else None
+    )
 
     sig = inspect.signature(create_fn)
     areal_client_ignored_args = ["model"] + (extra_ignored_args or [])
@@ -621,6 +633,22 @@ def _is_default_value(k: str, v: Any) -> bool:
         kwargs["top_p"] = 1.0
         _warn_once("top_p not set in request, defaulting to 1.0")
 
+    if (
+        _deterministic_sampling
+        and kwargs.get("seed") is None
+        and "seed" in areal_client_allowed_args
+    ):
+        assert request_index is not None
+        # The logical identity excludes the physical session collision suffix.
+        # Reserve request indices at ingress so concurrent requests remain
+        # distinct without holding a lock during inference.
+        # TODO(agent): Strict mapping of concurrent sibling requests to seeds
+        # requires a stable caller-provided request identity. Group samples use
+        # separate sessions, so their sample_idx-based identities are stable.
+        kwargs["seed"] = _deterministic_sampling_seed(
+            session_data.sampling_seed_identity, request_index
+        )
+
     # Strip stream from request body to prevent it from bypassing the explicit
     # `stream` parameter.  Without this, a request with {"stream": true} would
     # leak through kwargs and cause the client to return an AsyncGenerator even

areal/experimental/openai/proxy/server.py

@@ -67,8 +67,14 @@ class ExportTrajectoriesResponse(BaseModel):
 class SessionData:
     """Data associated with a single RL session."""
 
-    def __init__(self, session_id: str, prefix_matcher=None):
+    def __init__(
+        self,
+        session_id: str,
+        prefix_matcher=None,
+        sampling_seed_identity: str | None = None,
+    ):
         self.session_id = session_id
+        self.sampling_seed_identity = sampling_seed_identity or session_id
 
         self._completed = False
         self._completions = InteractionCache(
@@ -80,6 +86,14 @@ def __init__(self, session_id: str, prefix_matcher=None):
         self._last_access_time = time.time()
         self._end_time = None
         self._lock = threading.Lock()
+        self._next_sampling_request_index = 0
+
+    def next_sampling_request_index(self) -> int:
+        """Reserve a unique request index without serializing request execution."""
+        with self._lock:
+            request_index = self._next_sampling_request_index
+            self._next_sampling_request_index += 1
+        return request_index
 
     def update_last_access(self):
         """Update the last access time for this session."""

areal/experimental/openai/proxy/workflow.py

@@ -166,7 +166,15 @@ async def _grant_capacity(self, session: aiohttp.ClientSession) -> None:
     async def arun_episode(
         self, engine: TRolloutEngine, data: dict[str, Any]
     ) -> dict[str, InteractionWithTokenLogpReward] | None:
-        task_id = workflow_context.get().task_id
+        context = workflow_context.get()
+        task_id = context.task_id
+        # Qualify the proxy session with the group sample index so each group
+        # member owns a distinct, run-stable session namespace.
+        proxy_task_id = (
+            f"{task_id}:{context.sample_idx}"
+            if context.sample_idx is not None
+            else str(task_id)
+        )
 
         http_session = await workflow_context.get_aiohttp_session()
 
@@ -190,7 +198,7 @@ async def arun_episode(
             proxy_client = OpenAIProxyClient(
                 session=http_session,
                 base_url=self.proxy_addr,
-                task_id=str(task_id),
+                task_id=proxy_task_id,
                 admin_api_key=self._admin_api_key,
             )
             proxy_client.session_id = session_info.session_id
@@ -220,7 +228,7 @@ async def arun_episode(
         proxy_client = OpenAIProxyClient(
             session=http_session,
             base_url=self.proxy_addr,
-            task_id=str(task_id),
+            task_id=proxy_task_id,
             admin_api_key=self._admin_api_key,
         )
         async with proxy_client:

areal/infra/controller/rollout_controller.py

@@ -222,6 +222,7 @@ def initialize(
             task_factory=self._create_submit_callback,
             staleness_manager=self._staleness_manager,
             enable_tracing=self.config.enable_rollout_tracing,
+            deterministic_order=getattr(self.config, "deterministic_sampling", False),
         )
         # Initialize the dispatcher's async task runner
         self._dispatcher.initialize(logger=logger)

areal/infra/remote_inf_engine.py

@@ -89,9 +89,32 @@ async def arun_episode(
     ) -> dict[str, Any] | None:
         from areal.experimental.openai import InteractionWithTokenLogpReward
 
-        results = await asyncio.gather(
-            *[self.workflow.arun_episode(engine, data) for _ in range(self.group_size)]
+        async def run_sample(sample_idx: int) -> tuple[int, Any]:
+            from areal.infra import workflow_context
+            from areal.infra.workflow_context import WorkflowContext
+
+            parent = workflow_context.get()
+            workflow_context.set(
+                WorkflowContext(
+                    is_eval=parent.is_eval,
+                    task_id=parent.task_id,
+                    sample_idx=sample_idx,
+                )
+            )
+            result = await self.workflow.arun_episode(engine, data)
+            return sample_idx, result
+
+        indexed_results = await asyncio.gather(
+            *[run_sample(sample_idx) for sample_idx in range(self.group_size)]
         )
+        indexed_results.sort(key=lambda item: item[0])
+        sample_indices = [sample_idx for sample_idx, _ in indexed_results]
+        if sample_indices != list(range(self.group_size)):
+            raise RuntimeError(
+                "Grouped rollout returned invalid sample indices: "
+                f"expected {list(range(self.group_size))}, got {sample_indices}"
+            )
+        results = [result for _, result in indexed_results]
 
         valid_results = [r for r in results if r is not None]
 

areal/infra/workflow_context.py

@@ -31,10 +31,15 @@ class WorkflowContext:
         Whether the workflow is running in evaluation mode.
     task_id : int | None
         The task ID assigned by the workflow executor.
+    sample_idx : int | None
+        Index of this sample within its rollout group, when the workflow runs
+        under a grouped workflow. Gives group members a stable identity that
+        does not depend on completion order.
     """
 
     is_eval: bool = False
     task_id: int | None = None
+    sample_idx: int | None = None
 
 
 _current_context: ContextVar[WorkflowContext] = ContextVar(