Auto-sync: 2026-06-29 08:12:42

dovla_cil/eval/maniskill_policy_rollout.py

@@ -22,7 +22,12 @@ from dovla_cil.models.dovla import (
 from dovla_cil.utils.io import read_json, write_json
 
 
-_NUMPY_RESIDUAL_REDUCERS = {"mean_by_type", "median_by_type", "kernel_mean_by_type"}
+_NUMPY_RESIDUAL_REDUCERS = {
+    "mean_by_type",
+    "median_by_type",
+    "kernel_mean_by_type",
+    "compose_mean_by_type",
+}
 _FIELD_CONDITIONED_RESIDUAL_REDUCERS = {"field_softmax"}
 _RESIDUAL_REDUCERS = (
     {"none"} | _NUMPY_RESIDUAL_REDUCERS | _FIELD_CONDITIONED_RESIDUAL_REDUCERS
@@ -194,7 +199,8 @@ def evaluate_maniskill_policy_rollout(
     if retrieval_residual_reduce not in _RESIDUAL_REDUCERS:
         raise ValueError(
             "retrieval_residual_reduce must be 'none', 'mean_by_type', "
-            "'median_by_type', 'kernel_mean_by_type', or 'field_softmax'"
+            "'median_by_type', 'kernel_mean_by_type', 'compose_mean_by_type', "
+            "or 'field_softmax'"
         )
     if not 0.0 <= retrieval_type_min_success <= 1.0:
         raise ValueError("retrieval_type_min_success must be in [0, 1]")
@@ -820,9 +826,15 @@ def _reduce_residual_candidates_by_type(
 ) -> tuple[list[list[list[float]]], list[str]] | tuple[
     list[list[list[float]]], list[str], list[float]
 ]:
-    if mode not in {"mean_by_type", "median_by_type", "kernel_mean_by_type"}:
+    if mode not in {
+        "mean_by_type",
+        "median_by_type",
+        "kernel_mean_by_type",
+        "compose_mean_by_type",
+    }:
         raise ValueError(
-            "mode must be 'mean_by_type', 'median_by_type', or 'kernel_mean_by_type'"
+            "mode must be 'mean_by_type', 'median_by_type', 'kernel_mean_by_type', "
+            "or 'compose_mean_by_type'"
         )
     if len(residuals) != len(candidate_types):
         raise ValueError("residuals and candidate_types must have the same length")
@@ -852,7 +864,7 @@ def _reduce_residual_candidates_by_type(
         if not values:
             continue
         stack = np.stack(values, axis=0)
-        if mode == "mean_by_type":
+        if mode in {"mean_by_type", "compose_mean_by_type"}:
             reduced = np.mean(stack, axis=0)
             reduced_bonus = float(np.mean(value_bonuses)) if value_bonuses else 0.0
         elif mode == "kernel_mean_by_type":
@@ -887,6 +899,23 @@ def _reduce_residual_candidates_by_type(
         reduced_residuals.append(reduced.astype(np.float32).tolist())
         reduced_types.append(candidate_type)
         reduced_bonuses.append(reduced_bonus)
+    if mode == "compose_mean_by_type":
+        base_count = len(reduced_residuals)
+        for left_index in range(base_count):
+            left_type = reduced_types[left_index]
+            if left_type == "policy_residual":
+                continue
+            left = np.asarray(reduced_residuals[left_index], dtype=np.float32)
+            for right_index in range(left_index + 1, base_count):
+                right_type = reduced_types[right_index]
+                if right_type == "policy_residual":
+                    continue
+                right = np.asarray(reduced_residuals[right_index], dtype=np.float32)
+                reduced_residuals.append((left + right).astype(np.float32).tolist())
+                reduced_types.append(f"{left_type}+{right_type}")
+                reduced_bonuses.append(
+                    0.5 * (reduced_bonuses[left_index] + reduced_bonuses[right_index])
+                )
     if bonuses is None:
         return reduced_residuals, reduced_types
     return reduced_residuals, reduced_types, reduced_bonuses