fix(rollout): train safely on incomplete groups

areal/api/cli_args.py

@@ -77,6 +77,11 @@ class NormConfig:
         default=1, metadata={"help": "Group size for group-level normalization"}
     )
 
+    @property
+    def uses_group_statistics(self) -> bool:
+        """Whether normalization derives statistics from prompt groups."""
+        return self.mean_level == "group" or self.std_level == "group"
+
     def __post_init__(self):
         """Validate normalization configuration."""
         valid_levels = {"batch", "group", None}
@@ -1737,6 +1742,18 @@ class PPOActorConfig(TrainEngineConfig):
         default=None, metadata={"help": "Normalization configuration for advantages."}
     )
 
+    # Partial rollout groups
+    min_usable_group_size: int | None = field(
+        default=None,
+        metadata={
+            "help": "Minimum usable rollout slots a prompt group must keep to stay "
+            "trainable when some slots fail or are filtered. None derives the "
+            "minimum from reward_norm/adv_norm: 2 when either uses group "
+            "statistics (1 for a singleton target group), else 1. Only the v1 "
+            "rollout path consumes this option."
+        },
+    )
+
     # KL Control
     kl_ctl: float = field(default=0.1, metadata={"help": "KL divergence coefficient"})
     kl_estimator: str = field(
@@ -1828,6 +1845,34 @@ class PPOActorConfig(TrainEngineConfig):
         metadata={"help": "Maximum number of new tokens to generate"},
     )
 
+    def _uses_group_statistics(self) -> bool:
+        for normalization in (self.reward_norm, self.adv_norm):
+            if normalization is None:
+                continue
+            if isinstance(normalization, (dict, DictConfig)):
+                if (
+                    normalization.get("mean_level") == "group"
+                    or normalization.get("std_level") == "group"
+                ):
+                    return True
+            elif normalization.uses_group_statistics:
+                return True
+        return False
+
+    def resolve_min_usable_group_size(self, target_group_size: int) -> int:
+        """Minimum usable rollout slots a group must keep to stay trainable.
+
+        An explicit ``min_usable_group_size`` wins. Otherwise group-relative
+        normalization needs at least two group members before partial groups
+        become a hazard; a singleton target group is complete by definition,
+        so it keeps the minimum of one.
+        """
+        if self.min_usable_group_size is not None:
+            return self.min_usable_group_size
+        if self._uses_group_statistics():
+            return min(2, target_group_size)
+        return 1
+
     def should_compute_prox_logp(self) -> bool:
         """Determine if forward pass is needed for proximal log-probabilities.
 
@@ -1859,6 +1904,20 @@ def __post_init__(self):
                 f"{self.gae_timestep_unit!r}"
             )
 
+        if self.min_usable_group_size is not None:
+            if self.min_usable_group_size < 1:
+                raise ValueError(
+                    "min_usable_group_size must be a positive integer, "
+                    f"got {self.min_usable_group_size}"
+                )
+            if self.min_usable_group_size < 2 and self._uses_group_statistics():
+                raise ValueError(
+                    "min_usable_group_size must be at least 2 when reward_norm or "
+                    "adv_norm uses group statistics: a lone surviving rollout has "
+                    "no group peers to normalize against. Leave it unset to derive "
+                    "the minimum instead."
+                )
+
         reward_norm = self.reward_norm
         if isinstance(reward_norm, (dict, DictConfig)):
             reward_mean_level = reward_norm.get("mean_level")

areal/api/engine_api.py

@@ -202,6 +202,7 @@ def rollout_batch(
         group_size: int = 1,
         reward_normalization: bool = False,
         drop_incomplete_group: bool = False,
+        min_usable_group_size: int = 1,
     ) -> list[dict[str, Any]]:
         """Submit a batch of requests and wait for results.
 
@@ -221,12 +222,16 @@ def rollout_batch(
         group_size : int, optional
             Number of times to run the workflow per input and concatenate results.
             Default is 1 (no grouping).
+        min_usable_group_size : int, optional
+            Estimator-owned minimum number of usable logical rollout slots. Must be
+            between 1 and ``group_size``. Default is 1.
 
         Returns
         -------
         list[dict[str, Any]]
             A list of trajectory dictionaries, one per accepted rollout result.
-            Each trajectory contains tensors with shape [group_size, seqlen, ...].
+            Each trajectory contains tensors whose leading dimension is the number
+            of usable slots, between ``min_usable_group_size`` and ``group_size``.
         """
         raise NotImplementedError()
 
@@ -241,6 +246,7 @@ def prepare_batch(
         dynamic_bs: bool = False,
         reward_normalization: bool = False,
         drop_incomplete_group: bool = False,
+        min_usable_group_size: int = 1,
     ) -> list[dict[str, Any]]:
         """Prepare a batch of data for training from a dataloader.
 
@@ -261,6 +267,9 @@ def prepare_batch(
             If True, enables dynamic batch sizing. The method will stop collecting
             when (accepted + rejected) >= batch_size, returning only accepted results.
             This results in variable-sized batches of valid data. Default is False.
+        min_usable_group_size : int, optional
+            Estimator-owned minimum number of usable logical rollout slots. Must be
+            between 1 and ``group_size``. Default is 1.
 
         Returns
         -------
@@ -741,6 +750,7 @@ def submit(
         is_eval: bool = False,
         reward_normalization: bool = False,
         drop_incomplete_group: bool = False,
+        min_usable_group_size: int = 1,
     ) -> int:
         """Submit a request to the inference engine and return immediately.
 
@@ -773,6 +783,9 @@ def submit(
         is_eval : bool, optional
             Whether this is an evaluation workflow. Affects variables like trajectory dump path
             and statistics keys. By default False.
+        min_usable_group_size : int, optional
+            Estimator-owned minimum number of usable logical rollout slots. Must be
+            between 1 and ``group_size``. Default is 1.
 
         Returns
         -------
@@ -846,6 +859,7 @@ def rollout_batch(
         group_size: int = 1,
         reward_normalization: bool = False,
         drop_incomplete_group: bool = False,
+        min_usable_group_size: int = 1,
     ) -> list[dict[str, Any]]:
         """Submit a batch of requests to the inference engine and wait for the results.
 
@@ -873,6 +887,9 @@ def rollout_batch(
         group_size : int, optional
             Number of times to run the workflow per input and concatenate results.
             Default is 1 (no grouping).
+        min_usable_group_size : int, optional
+            Estimator-owned minimum number of usable logical rollout slots. Must be
+            between 1 and ``group_size``. Default is 1.
 
         Returns
         -------
@@ -893,6 +910,7 @@ def prepare_batch(
         dynamic_bs: bool = False,
         reward_normalization: bool = False,
         drop_incomplete_group: bool = False,
+        min_usable_group_size: int = 1,
     ) -> list[dict[str, Any]]:
         """Asynchronously submit and wait until a full batch is ready with controlled staleness.
 
@@ -902,9 +920,10 @@ def prepare_batch(
 
             This method caches an internal data generator on the first call.
             The ``dataloader``, ``workflow``, ``workflow_kwargs``, ``group_size``,
-            and ``should_accept_fn`` parameters are captured at the first invocation
-            and reused in all subsequent calls. Passing different arguments in
-            later calls will **not** take effect.
+            ``reward_normalization``, ``drop_incomplete_group``,
+            ``min_usable_group_size``, and ``should_accept_fn`` parameters are captured
+            at the first invocation and reused in all subsequent calls. Passing
+            different arguments in later calls will **not** take effect.
 
             If you need to switch configurations mid-training, consider:
 
@@ -936,6 +955,9 @@ def prepare_batch(
             If True, enables dynamic batch sizing. The method will stop collecting
             when (accepted + rejected) >= batch_size, returning only accepted results.
             This results in variable-sized batches of valid data. Default is False.
+        min_usable_group_size : int, optional
+            Estimator-owned minimum number of usable logical rollout slots. Must be
+            between 1 and ``group_size``. Default is 1.
 
         Returns
         -------

areal/api/workflow_api.py

@@ -22,6 +22,22 @@ async def arun_episode(
         ----
         Returning `None` implies that this trajectory is rejected and will not be used for training.
 
+        When the resolved ``min_usable_group_size`` is at least 2 — derived
+        from group-relative normalization (``mean_level='group'`` or
+        ``std_level='group'`` in ``actor.reward_norm`` / ``actor.adv_norm``),
+        or set explicitly via ``actor.min_usable_group_size`` — and rollouts
+        are grouped (``n_samples >= 2``), each ``arun_episode`` call must contribute
+        exactly one training sample: a tensor dict with batch size 1, or a
+        dict containing a single interaction. Returning multiple samples per
+        episode (e.g. one row per turn, or tree-search branches) raises a
+        non-retryable ``WorkflowContractError`` that terminates training,
+        because group statistics would otherwise treat same-episode rows as
+        independent group members. Ungrouped rollouts (``n_samples=1``)
+        install no group wrapper and are not checked; a multi-sample episode
+        is then normalized as its own group of same-episode rows. To train
+        multi-sample workflows, use batch-level normalization or merge each
+        episode into a single sequence.
+
         See concrete example implementations under the `areal/workflow` directory.
 
         Parameters

areal/engine/core/__init__.py

@@ -4,12 +4,14 @@
 
 from areal.engine.core.train_engine import (
     aggregate_eval_losses,
+    compute_microbatch_loss_weight,
     compute_total_loss_weight,
     reorder_and_pad_outputs,
 )
 
 __all__ = [
     "aggregate_eval_losses",
+    "compute_microbatch_loss_weight",
     "compute_total_loss_weight",
     "reorder_and_pad_outputs",
 ]

areal/engine/core/train_engine.py

@@ -14,19 +14,37 @@
 
 from areal.infra.platforms import current_platform
 from areal.utils.data import (
+    TRANSPORT_DUMMY_KEY,
     MicroBatchList,
     pad_and_stack_tensors_along_first_dim,
     reorder_list,
     unpack_sequence,
 )
 
 __all__ = [
+    "compute_microbatch_loss_weight",
     "compute_total_loss_weight",
     "aggregate_eval_losses",
     "reorder_and_pad_outputs",
 ]
 
 
+def compute_microbatch_loss_weight(
+    microbatch: dict[str, Any],
+    loss_weight_fn: Callable[[dict[str, Any]], torch.Tensor],
+) -> torch.Tensor:
+    """Return zero without invoking an objective on transport-only data."""
+    if microbatch.get(TRANSPORT_DUMMY_KEY) is not True:
+        return loss_weight_fn(microbatch)
+    reference = next(
+        (value for value in microbatch.values() if isinstance(value, torch.Tensor)),
+        None,
+    )
+    if reference is None:
+        raise ValueError("Transport micro-batch does not contain a tensor")
+    return torch.zeros((), dtype=torch.float32, device=reference.device)
+
+
 def compute_total_loss_weight(
     mb_list: MicroBatchList,
     loss_weight_fn: Callable[[dict[str, Any]], torch.Tensor],
@@ -52,7 +70,9 @@ def compute_total_loss_weight(
         The total loss weight (scalar tensor) after all_reduce.
     """
     total_weight = (
-        torch.stack([loss_weight_fn(mb) for mb in mb_list.mbs])
+        torch.stack(
+            [compute_microbatch_loss_weight(mb, loss_weight_fn) for mb in mb_list.mbs]
+        )
         .sum()
         .detach()
         .clone()
@@ -138,7 +158,11 @@ def reorder_and_pad_outputs(
         The processed outputs, padded and stacked along batch dimension.
     """
     res = aggregate_fn(outputs)
+    semantic_batch_size = len(output_seqlens)
+    output_seqlens = [*output_seqlens, *([1] * mb_list.transport_dummy_count)]
     seqlens = [output_seqlens[i] for i in mb_list.forward_indices]
     unpacked = unpack_sequence(res, lens=seqlens, dim=0)
     reordered = reorder_list(unpacked, mb_list.backward_indices)
+    if mb_list.transport_dummy_count:
+        reordered = reordered[:semantic_batch_size]
     return pad_and_stack_tensors_along_first_dim(reordered)