feat(reward): preserve episode outcomes and support mean-only groups

areal/api/cli_args.py

@@ -260,6 +260,13 @@ class GenerationHyperparameters:
             )
         },
     )
+    reward_normalization_use_std: bool = field(
+        default=True,
+        metadata={
+            "help": "Divide grouped rollout rewards by their standard deviation. "
+            "Set False to subtract only the mean when reward_normalization is enabled."
+        },
+    )
     drop_incomplete_group: bool = field(
         default=False,
         metadata={
@@ -324,6 +331,7 @@ def to_openai_agents_model_settings_dict(
     # OpenAI client kwargs even when users enable them.
     _WORKFLOW_ONLY_ARGS: ClassVar[set[str]] = {
         "reward_normalization",
+        "reward_normalization_use_std",
         "drop_incomplete_group",
     }
 

areal/api/engine_api.py

@@ -203,6 +203,7 @@ def rollout_batch(
         reward_normalization: bool = False,
         drop_incomplete_group: bool = False,
         min_usable_group_size: int = 1,
+        reward_normalization_use_std: bool = True,
     ) -> list[dict[str, Any]]:
         """Submit a batch of requests and wait for results.
 
@@ -247,6 +248,7 @@ def prepare_batch(
         reward_normalization: bool = False,
         drop_incomplete_group: bool = False,
         min_usable_group_size: int = 1,
+        reward_normalization_use_std: bool = True,
     ) -> list[dict[str, Any]]:
         """Prepare a batch of data for training from a dataloader.
 
@@ -751,6 +753,7 @@ def submit(
         reward_normalization: bool = False,
         drop_incomplete_group: bool = False,
         min_usable_group_size: int = 1,
+        reward_normalization_use_std: bool = True,
     ) -> int:
         """Submit a request to the inference engine and return immediately.
 
@@ -860,6 +863,7 @@ def rollout_batch(
         reward_normalization: bool = False,
         drop_incomplete_group: bool = False,
         min_usable_group_size: int = 1,
+        reward_normalization_use_std: bool = True,
     ) -> list[dict[str, Any]]:
         """Submit a batch of requests to the inference engine and wait for the results.
 
@@ -911,6 +915,7 @@ def prepare_batch(
         reward_normalization: bool = False,
         drop_incomplete_group: bool = False,
         min_usable_group_size: int = 1,
+        reward_normalization_use_std: bool = True,
     ) -> list[dict[str, Any]]:
         """Asynchronously submit and wait until a full batch is ready with controlled staleness.
 

areal/engine/fsdp_engine.py

@@ -586,6 +586,7 @@ def rollout_batch(
         reward_normalization: bool = False,
         drop_incomplete_group: bool = False,
         min_usable_group_size: int = 1,
+        reward_normalization_use_std: bool = True,
     ) -> list[dict[str, Any]]:
         self._check_rollout_engine_connected()
         return self.rollout_coordinator.rollout_batch(
@@ -595,6 +596,7 @@ def rollout_batch(
             group_size=group_size,
             min_usable_group_size=min_usable_group_size,
             reward_normalization=reward_normalization,
+            reward_normalization_use_std=reward_normalization_use_std,
             drop_incomplete_group=drop_incomplete_group,
         )
 
@@ -609,6 +611,7 @@ def prepare_batch(
         reward_normalization: bool = False,
         drop_incomplete_group: bool = False,
         min_usable_group_size: int = 1,
+        reward_normalization_use_std: bool = True,
     ) -> list[dict[str, Any]]:
         self._check_rollout_engine_connected()
         return self.rollout_coordinator.prepare_batch(
@@ -620,6 +623,7 @@ def prepare_batch(
             min_usable_group_size=min_usable_group_size,
             dynamic_bs=dynamic_bs,
             reward_normalization=reward_normalization,
+            reward_normalization_use_std=reward_normalization_use_std,
             drop_incomplete_group=drop_incomplete_group,
         )
 

areal/engine/megatron_engine.py

@@ -1040,6 +1040,7 @@ def rollout_batch(
         reward_normalization: bool = False,
         drop_incomplete_group: bool = False,
         min_usable_group_size: int = 1,
+        reward_normalization_use_std: bool = True,
     ) -> list[dict[str, Any]]:
         self._check_rollout_engine_connected()
         return self.rollout_coordinator.rollout_batch(
@@ -1049,6 +1050,7 @@ def rollout_batch(
             group_size=group_size,
             min_usable_group_size=min_usable_group_size,
             reward_normalization=reward_normalization,
+            reward_normalization_use_std=reward_normalization_use_std,
             drop_incomplete_group=drop_incomplete_group,
         )
 
@@ -1063,6 +1065,7 @@ def prepare_batch(
         reward_normalization: bool = False,
         drop_incomplete_group: bool = False,
         min_usable_group_size: int = 1,
+        reward_normalization_use_std: bool = True,
     ) -> list[dict[str, Any]]:
         self._check_rollout_engine_connected()
         return self.rollout_coordinator.prepare_batch(
@@ -1074,6 +1077,7 @@ def prepare_batch(
             min_usable_group_size=min_usable_group_size,
             dynamic_bs=dynamic_bs,
             reward_normalization=reward_normalization,
+            reward_normalization_use_std=reward_normalization_use_std,
             drop_incomplete_group=drop_incomplete_group,
         )
 

areal/engine/sglang_remote.py

@@ -676,6 +676,7 @@ def submit(
         reward_normalization: bool = False,
         drop_incomplete_group: bool = False,
         min_usable_group_size: int = 1,
+        reward_normalization_use_std: bool = True,
     ) -> int:
         """Submit a request to the inference engine."""
         return self._engine.submit(
@@ -690,6 +691,7 @@ def submit(
             is_eval=is_eval,
             proxy_addr=proxy_addr,
             reward_normalization=reward_normalization,
+            reward_normalization_use_std=reward_normalization_use_std,
             drop_incomplete_group=drop_incomplete_group,
         )
 
@@ -719,6 +721,7 @@ def rollout_batch(
         reward_normalization: bool = False,
         drop_incomplete_group: bool = False,
         min_usable_group_size: int = 1,
+        reward_normalization_use_std: bool = True,
     ) -> dict[str, Any]:
         """Submit a batch of requests and wait for results.
 
@@ -732,6 +735,7 @@ def rollout_batch(
             group_size=group_size,
             min_usable_group_size=min_usable_group_size,
             reward_normalization=reward_normalization,
+            reward_normalization_use_std=reward_normalization_use_std,
             drop_incomplete_group=drop_incomplete_group,
         )
 
@@ -746,6 +750,7 @@ def prepare_batch(
         reward_normalization: bool = False,
         drop_incomplete_group: bool = False,
         min_usable_group_size: int = 1,
+        reward_normalization_use_std: bool = True,
     ):
         """Asynchronously submit and wait until a full batch is ready."""
         return self._engine.prepare_batch(
@@ -757,6 +762,7 @@ def prepare_batch(
             min_usable_group_size=min_usable_group_size,
             dynamic_bs=dynamic_bs,
             reward_normalization=reward_normalization,
+            reward_normalization_use_std=reward_normalization_use_std,
             drop_incomplete_group=drop_incomplete_group,
         )
 

areal/engine/vllm_remote.py

@@ -440,6 +440,7 @@ def submit(
         reward_normalization: bool = False,
         drop_incomplete_group: bool = False,
         min_usable_group_size: int = 1,
+        reward_normalization_use_std: bool = True,
     ) -> int:
         """Submit a request to the inference engine."""
         return self._engine.submit(
@@ -454,6 +455,7 @@ def submit(
             is_eval=is_eval,
             proxy_addr=proxy_addr,
             reward_normalization=reward_normalization,
+            reward_normalization_use_std=reward_normalization_use_std,
             drop_incomplete_group=drop_incomplete_group,
         )
 
@@ -483,6 +485,7 @@ def rollout_batch(
         reward_normalization: bool = False,
         drop_incomplete_group: bool = False,
         min_usable_group_size: int = 1,
+        reward_normalization_use_std: bool = True,
     ) -> dict[str, Any]:
         """Submit a batch of requests and wait for results.
 
@@ -496,6 +499,7 @@ def rollout_batch(
             group_size=group_size,
             min_usable_group_size=min_usable_group_size,
             reward_normalization=reward_normalization,
+            reward_normalization_use_std=reward_normalization_use_std,
             drop_incomplete_group=drop_incomplete_group,
         )
 
@@ -510,6 +514,7 @@ def prepare_batch(
         reward_normalization: bool = False,
         drop_incomplete_group: bool = False,
         min_usable_group_size: int = 1,
+        reward_normalization_use_std: bool = True,
     ):
         """Asynchronously submit and wait until a full batch is ready."""
         return self._engine.prepare_batch(
@@ -521,6 +526,7 @@ def prepare_batch(
             min_usable_group_size=min_usable_group_size,
             dynamic_bs=dynamic_bs,
             reward_normalization=reward_normalization,
+            reward_normalization_use_std=reward_normalization_use_std,
             drop_incomplete_group=drop_incomplete_group,
         )
 

areal/experimental/engine/archon_engine.py

@@ -700,6 +700,7 @@ def rollout_batch(
         reward_normalization: bool = False,
         drop_incomplete_group: bool = False,
         min_usable_group_size: int = 1,
+        reward_normalization_use_std: bool = True,
     ) -> list[dict[str, Any]]:
         """Perform rollout using connected inference engine."""
         self._check_rollout_engine_connected()
@@ -710,6 +711,7 @@ def rollout_batch(
             group_size=group_size,
             min_usable_group_size=min_usable_group_size,
             reward_normalization=reward_normalization,
+            reward_normalization_use_std=reward_normalization_use_std,
             drop_incomplete_group=drop_incomplete_group,
         )
 
@@ -724,6 +726,7 @@ def prepare_batch(
         reward_normalization: bool = False,
         drop_incomplete_group: bool = False,
         min_usable_group_size: int = 1,
+        reward_normalization_use_std: bool = True,
     ) -> list[dict[str, Any]]:
         """Prepare batch from dataloader with rollout."""
         self._check_rollout_engine_connected()
@@ -736,6 +739,7 @@ def prepare_batch(
             min_usable_group_size=min_usable_group_size,
             dynamic_bs=dynamic_bs,
             reward_normalization=reward_normalization,
+            reward_normalization_use_std=reward_normalization_use_std,
             drop_incomplete_group=drop_incomplete_group,
         )
 

areal/experimental/openai/types.py

@@ -357,6 +357,8 @@ def normalize_group_rewards(
 
 def normalize_logical_rollout_rewards(
     results: list[dict[str, InteractionWithTokenLogpReward] | None],
+    *,
+    use_std: bool = True,
 ) -> bool:
     """Normalize row rewards against one explicit or equal-row rollout reference."""
     if not results or any(not result for result in results):
@@ -390,7 +392,7 @@ def normalize_logical_rollout_rewards(
     std = (
         rewards.std(unbiased=False) if len(references) > 1 else rewards.new_tensor(1.0)
     )
-    scale = torch.where(std <= 1e-8, 1.0, std + 1e-8)
+    scale = torch.where(std <= 1e-8, 1.0, std + 1e-8) if use_std else 1.0
     row_rewards = torch.tensor(
         [v.reward for result in results for v in result.values()], dtype=torch.float32
     )

areal/infra/controller/rollout_controller.py

@@ -156,6 +156,7 @@ class _RemoteRolloutTaskInput:
     min_usable_group_size: int = 1
     proxy_addr: str | None = None
     reward_normalization: bool = False
+    reward_normalization_use_std: bool = True
     drop_incomplete_group: bool = False
 
 
@@ -1098,6 +1099,7 @@ async def _submit_then_wait() -> WorkflowTaskResult | None:
                     callback_addr=f"http://{self.callback_addr}/callback/rollout_complete",
                     proxy_addr=proxy_addr,
                     reward_normalization=pending_task.reward_normalization,
+                    reward_normalization_use_std=pending_task.reward_normalization_use_std,
                     drop_incomplete_group=pending_task.drop_incomplete_group,
                 )
 
@@ -1169,6 +1171,7 @@ def submit(
         reward_normalization: bool = False,
         drop_incomplete_group: bool = False,
         min_usable_group_size: int = 1,
+        reward_normalization_use_std: bool = True,
     ) -> int:
         validate_rollout_group_sizes(group_size, min_usable_group_size)
 
@@ -1197,6 +1200,7 @@ def submit(
             min_usable_group_size=min_usable_group_size,
             proxy_addr=proxy_addr,
             reward_normalization=reward_normalization,
+            reward_normalization_use_std=reward_normalization_use_std,
             drop_incomplete_group=drop_incomplete_group,
         )
 
@@ -1229,6 +1233,7 @@ def rollout_batch(
         reward_normalization: bool = False,
         drop_incomplete_group: bool = False,
         min_usable_group_size: int = 1,
+        reward_normalization_use_std: bool = True,
     ) -> list[dict[str, Any]]:
         perf_tracer.instant(
             "rollout_controller.rollout_batch",
@@ -1244,6 +1249,7 @@ def rollout_batch(
                 group_size=group_size,
                 min_usable_group_size=min_usable_group_size,
                 reward_normalization=reward_normalization,
+                reward_normalization_use_std=reward_normalization_use_std,
                 drop_incomplete_group=drop_incomplete_group,
             )
         results = self.wait(count=len(data))
@@ -1262,6 +1268,7 @@ def prepare_batch(
         reward_normalization: bool = False,
         drop_incomplete_group: bool = False,
         min_usable_group_size: int = 1,
+        reward_normalization_use_std: bool = True,
     ) -> list[dict[str, Any]]:
         """Prepare a batch with controlled staleness.
 
@@ -1292,6 +1299,7 @@ def task_input_generator():
                         group_size=group_size,
                         min_usable_group_size=min_usable_group_size,
                         reward_normalization=reward_normalization,
+                        reward_normalization_use_std=reward_normalization_use_std,
                         drop_incomplete_group=drop_incomplete_group,
                     )
 

areal/infra/controller/train_controller.py

@@ -852,6 +852,7 @@ def prepare_batch(
         reward_normalization: bool = False,
         drop_incomplete_group: bool = False,
         min_usable_group_size: int = 1,
+        reward_normalization_use_std: bool = True,
     ) -> list[dict[str, Any]]:
         return self.rollout.prepare_batch(
             dataloader=dataloader,
@@ -862,6 +863,7 @@ def prepare_batch(
             min_usable_group_size=min_usable_group_size,
             dynamic_bs=dynamic_bs,
             reward_normalization=reward_normalization,
+            reward_normalization_use_std=reward_normalization_use_std,
             drop_incomplete_group=drop_incomplete_group,
         )
 
@@ -875,6 +877,7 @@ def rollout_batch(
         reward_normalization: bool = False,
         drop_incomplete_group: bool = False,
         min_usable_group_size: int = 1,
+        reward_normalization_use_std: bool = True,
     ) -> list[dict[str, Any]]:
         return self.rollout.rollout_batch(
             data=data,
@@ -884,6 +887,7 @@ def rollout_batch(
             group_size=group_size,
             min_usable_group_size=min_usable_group_size,
             reward_normalization=reward_normalization,
+            reward_normalization_use_std=reward_normalization_use_std,
             drop_incomplete_group=drop_incomplete_group,
         )
 

areal/infra/dist_rollout.py

@@ -269,6 +269,7 @@ def rollout_batch(
         reward_normalization: bool = False,
         drop_incomplete_group: bool = False,
         min_usable_group_size: int = 1,
+        reward_normalization_use_std: bool = True,
     ) -> list[dict[str, Any]]:
         """Generate rollout batch with distributed coordination (synchronous).
 
@@ -311,6 +312,7 @@ def rollout_batch(
                 group_size=group_size,
                 min_usable_group_size=min_usable_group_size,
                 reward_normalization=reward_normalization,
+                reward_normalization_use_std=reward_normalization_use_std,
                 drop_incomplete_group=drop_incomplete_group,
             )
         )
@@ -330,6 +332,7 @@ def prepare_batch(
         reward_normalization: bool = False,
         drop_incomplete_group: bool = False,
         min_usable_group_size: int = 1,
+        reward_normalization_use_std: bool = True,
     ) -> list[dict[str, Any]]:
         """Prepare async rollout batch with distributed coordination.
 
@@ -388,6 +391,7 @@ def _prepare_until_dispatchable() -> list[dict[str, Any]]:
                         min_usable_group_size=min_usable_group_size,
                         dynamic_bs=dynamic_bs,
                         reward_normalization=reward_normalization,
+                        reward_normalization_use_std=reward_normalization_use_std,
                         drop_incomplete_group=drop_incomplete_group,
                     )
                 except Exception as exc: