Auto-sync: 2026-06-29 17:58:49

dovla_cil/eval/maniskill_policy_rollout.py

@@ -103,6 +103,7 @@ def evaluate_maniskill_policy_rollout(
     lattice_exclude_types: tuple[str, ...] = (),
     candidate_type_bonuses: dict[str, float] | None = None,
     candidate_type_bonus_components: bool = False,
+    candidate_oracle_rollouts: int = 0,
 ) -> dict[str, Any]:
     """Execute a checkpoint policy from restored ManiSkill CIL states.
 
@@ -187,6 +188,8 @@ def evaluate_maniskill_policy_rollout(
         raise ValueError("field_optim_l2_penalty must be non-negative")
     if selection_margin < 0:
         raise ValueError("selection_margin must be non-negative")
+    if candidate_oracle_rollouts < 0:
+        raise ValueError("candidate_oracle_rollouts must be non-negative")
     if retrieval_neighbors <= 0:
         raise ValueError("retrieval_neighbors must be positive")
     if retrieval_metric not in _RETRIEVAL_METRICS:
@@ -352,6 +355,7 @@ def evaluate_maniskill_policy_rollout(
                 lattice_exclude_types=lattice_exclude_types,
                 candidate_type_bonuses=candidate_type_bonuses,
                 candidate_type_bonus_components=candidate_type_bonus_components,
+                candidate_oracle_rollouts=candidate_oracle_rollouts,
             )
             rows.extend(task_rows)
             task_summaries[task_id] = _summarize_rows(task_rows)
@@ -465,6 +469,7 @@ def evaluate_maniskill_policy_rollout(
         "lattice_exclude_types": list(lattice_exclude_types),
         "candidate_type_bonuses": candidate_type_bonuses,
         "candidate_type_bonus_components": bool(candidate_type_bonus_components),
+        "candidate_oracle_rollouts": int(candidate_oracle_rollouts),
         "policy_rollout_success_rate": _mean([row["success"] for row in rows]),
         "policy_rollout_progress": _mean([row["progress"] for row in rows]),
         "oracle_success_rate": _mean([row["oracle_success"] for row in rows]),
@@ -490,6 +495,11 @@ def evaluate_maniskill_policy_rollout(
         ),
         "rows": rows,
     }
+    candidate_oracle_rows = [
+        row for row in rows if row.get("candidate_oracle_success") is not None
+    ]
+    if candidate_oracle_rows:
+        result.update(_candidate_oracle_summary(candidate_oracle_rows))
     if output_path is not None:
         write_json(result, output_path)
     return result
@@ -1141,6 +1151,7 @@ def _evaluate_task_cases(
     lattice_exclude_types: tuple[str, ...] = (),
     candidate_type_bonuses: dict[str, float] | None = None,
     candidate_type_bonus_components: bool = False,
+    candidate_oracle_rollouts: int = 0,
 ) -> list[dict[str, Any]]:
     rows: list[dict[str, Any]] = []
     for start in range(0, len(cases), group_batch_size):
@@ -1272,6 +1283,38 @@ def _evaluate_task_cases(
                 device=env_device,
             )
             del after_state
+            candidate_oracle = None
+            if candidate_oracle_rollouts > 0:
+                candidate_oracle = _execute_candidate_oracle_prefix(
+                    task_id,
+                    batch,
+                    model=model,
+                    model_config=model_config,
+                    gym=gym,
+                    env=env,
+                    torch=torch,
+                    device=device,
+                    env_device=env_device,
+                    env_kwargs=env_kwargs,
+                    env_action_dim=env_dim,
+                    observations=observations,
+                    instructions=[case.instruction for case in batch],
+                    selected_indices=candidate_index,
+                    selection_mode=selection_mode,
+                    candidate_oracle_rollouts=candidate_oracle_rollouts,
+                    num_candidates=num_candidates,
+                    candidate_sigma=candidate_sigma,
+                    selection_seed=selection_seed + start,
+                    retrieval_residual_scale=retrieval_residual_scale,
+                    retrieval_residual_scales=retrieval_residual_scales,
+                    retrieval_residual_reduce=retrieval_residual_reduce,
+                    retrieval_residual_action_l2_penalty=retrieval_residual_action_l2_penalty,
+                    action_low=action_low,
+                    action_high=action_high,
+                    lattice_exclude_types=lattice_exclude_types,
+                    candidate_type_bonuses=candidate_type_bonuses or {},
+                    candidate_type_bonus_components=candidate_type_bonus_components,
+                )
             for index, case in enumerate(batch):
                 progress = float(max(0.0, min(1.0, rewards[index, 0])))
                 success = bool(successes[index, 0])
@@ -1333,11 +1376,267 @@ def _evaluate_task_cases(
                         "candidate_source_group_id": case.candidate_source_group_id,
                     }
                 )
+                if candidate_oracle is not None:
+                    rows[-1].update(candidate_oracle[index])
         finally:
             env.close()
     return rows
 
 
+def _execute_candidate_oracle_prefix(
+    task_id: str,
+    batch: list[_RolloutCase],
+    *,
+    model: DoVLAModel,
+    model_config: DoVLAConfig,
+    gym: Any,
+    env: Any,
+    torch: Any,
+    device: str,
+    env_device: Any,
+    env_kwargs: dict[str, Any],
+    env_action_dim: int,
+    observations: Any,
+    instructions: list[str],
+    selected_indices: np.ndarray,
+    selection_mode: str,
+    candidate_oracle_rollouts: int,
+    num_candidates: int,
+    candidate_sigma: float,
+    selection_seed: int,
+    retrieval_residual_scale: float,
+    retrieval_residual_scales: tuple[float, ...],
+    retrieval_residual_reduce: str,
+    retrieval_residual_action_l2_penalty: float,
+    action_low: Any | None,
+    action_high: Any | None,
+    lattice_exclude_types: tuple[str, ...],
+    candidate_type_bonuses: dict[str, float],
+    candidate_type_bonus_components: bool,
+) -> list[dict[str, Any]]:
+    del model_config, env_device
+    if selection_mode != "retrieval_residual":
+        raise ValueError("candidate_oracle_rollouts currently supports retrieval_residual")
+    if retrieval_residual_reduce in _FIELD_CONDITIONED_RESIDUAL_REDUCERS:
+        raise ValueError(
+            "candidate_oracle_rollouts requires an explicit residual lattice, not field_softmax"
+        )
+    if candidate_oracle_rollouts <= 0:
+        return [{} for _ in batch]
+
+    action_residuals = torch.tensor(
+        [case.candidate_action_values for case in batch],
+        dtype=torch.float32,
+        device=device,
+    )
+    policy_mean = model.forward_policy(observations, instructions)
+    candidate_mask = (
+        _lattice_candidate_mask(
+            batch,
+            torch=torch,
+            device=device,
+            exclude_types=lattice_exclude_types,
+        )
+        if lattice_exclude_types
+        else None
+    )
+    candidate_type_bonus = (
+        _lattice_candidate_type_bonus(
+            batch,
+            torch=torch,
+            device=device,
+            candidate_type_bonuses=candidate_type_bonuses,
+            use_components=candidate_type_bonus_components,
+        )
+        if candidate_type_bonuses or any(case.candidate_score_bonuses for case in batch)
+        else None
+    )
+    candidates, candidate_mask, candidate_type_bonus = _build_residual_lattice_candidates(
+        policy_mean,
+        action_residuals.to(device=policy_mean.device, dtype=policy_mean.dtype),
+        torch=torch,
+        residual_scale=retrieval_residual_scale,
+        residual_scales=retrieval_residual_scales,
+        candidate_mask=candidate_mask,
+        candidate_type_bonus=candidate_type_bonus,
+        num_gaussian_candidates=num_candidates,
+        candidate_sigma=candidate_sigma,
+        selection_seed=selection_seed,
+    )
+    scored_candidates, potentials = _score_lattice_action_chunks(
+        model,
+        observations,
+        instructions,
+        candidates,
+        torch=torch,
+        action_low=(
+            action_low.unsqueeze(1)
+            if action_low is not None and action_low.ndim == 3
+            else action_low
+        ),
+        action_high=(
+            action_high.unsqueeze(1)
+            if action_high is not None and action_high.ndim == 3
+            else action_high
+        ),
+        candidate_mask=candidate_mask,
+        candidate_type_bonus=candidate_type_bonus,
+        action_l2_penalty=retrieval_residual_action_l2_penalty,
+        action_l2_penalty_base=policy_mean,
+    )
+    candidate_indices = _diagnostic_candidate_indices(
+        potentials,
+        selected_indices,
+        candidate_oracle_rollouts=candidate_oracle_rollouts,
+        torch=torch,
+    )
+    batch_index = torch.arange(len(batch), device=scored_candidates.device).reshape(-1, 1)
+    diagnostic_candidates = scored_candidates[batch_index, candidate_indices]
+    diagnostic_np = diagnostic_candidates.detach().cpu().numpy().astype(np.float32)
+    branch_count = int(diagnostic_np.shape[1])
+    flat = diagnostic_np.reshape(
+        len(batch) * branch_count,
+        diagnostic_np.shape[-2],
+        diagnostic_np.shape[-1],
+    )
+    flat = _adapt_action_dim(flat, env_action_dim)
+    flat = _clip_to_action_space(flat, env)
+    adapted = flat.reshape(len(batch), branch_count, flat.shape[-2], flat.shape[-1])
+
+    diag_env_kwargs = dict(env_kwargs)
+    diag_env_kwargs["num_envs"] = len(batch) * branch_count
+    diag_env = gym.make(task_id, **diag_env_kwargs)
+    diag_base_env = diag_env.unwrapped
+    try:
+        diag_env_device = getattr(diag_base_env, "device", torch.device(device))
+        after_state, rewards, successes, restore_error = execute_grouped_action_lattice_batch(
+            diag_base_env,
+            [case.state for case in batch],
+            adapted,
+            torch=torch,
+            device=diag_env_device,
+        )
+        del after_state
+    finally:
+        diag_env.close()
+
+    clipped_rewards = np.clip(rewards.astype(np.float32), 0.0, 1.0)
+    scores = clipped_rewards + successes.astype(np.float32)
+    best_branch = np.argmax(scores, axis=1)
+    indices_np = candidate_indices.detach().cpu().numpy()
+    potentials_np = potentials.gather(1, candidate_indices).detach().cpu().numpy()
+    rows: list[dict[str, Any]] = []
+    for row_index, case in enumerate(batch):
+        branch = int(best_branch[row_index])
+        candidate_index = int(indices_np[row_index, branch])
+        candidate_type = _selected_candidate_type(
+            case,
+            selected_index=candidate_index,
+            selection_mode=selection_mode,
+            residual_scale_count=_residual_scale_count(retrieval_residual_scales),
+            retrieval_residual_reduce=retrieval_residual_reduce,
+        )
+        rows.append(
+            {
+                "candidate_oracle_rollout_count": branch_count,
+                "candidate_oracle_success": bool(successes[row_index, branch]),
+                "candidate_oracle_progress": float(clipped_rewards[row_index, branch]),
+                "candidate_oracle_score": float(scores[row_index, branch]),
+                "candidate_oracle_regret": max(
+                    0.0,
+                    case.oracle_score - float(scores[row_index, branch]),
+                ),
+                "candidate_oracle_score_gain_over_selected": float(
+                    scores[row_index, branch] - scores[row_index, 0]
+                ),
+                "candidate_oracle_improves_selected": bool(
+                    scores[row_index, branch] > scores[row_index, 0] + 1.0e-9
+                ),
+                "candidate_oracle_candidate_index": candidate_index,
+                "candidate_oracle_candidate_type": candidate_type,
+                "candidate_oracle_residual_scale": _selected_residual_scale(
+                    case,
+                    selected_index=candidate_index,
+                    selection_mode=selection_mode,
+                    residual_scale=retrieval_residual_scale,
+                    residual_scales=retrieval_residual_scales,
+                    retrieval_residual_reduce=retrieval_residual_reduce,
+                ),
+                "candidate_oracle_field_potential": float(
+                    potentials_np[row_index, branch]
+                ),
+                "candidate_oracle_restore_error": float(restore_error),
+                "candidate_oracle_selected_branch_success": bool(successes[row_index, 0]),
+                "candidate_oracle_selected_branch_progress": float(
+                    clipped_rewards[row_index, 0]
+                ),
+                "candidate_oracle_selected_branch_score": float(scores[row_index, 0]),
+                "candidate_oracle_indices": [
+                    int(value) for value in indices_np[row_index].tolist()
+                ],
+                "candidate_oracle_types": [
+                    _selected_candidate_type(
+                        case,
+                        selected_index=int(value),
+                        selection_mode=selection_mode,
+                        residual_scale_count=_residual_scale_count(
+                            retrieval_residual_scales
+                        ),
+                        retrieval_residual_reduce=retrieval_residual_reduce,
+                    )
+                    for value in indices_np[row_index].tolist()
+                ],
+            }
+        )
+    return rows
+
+
+def _diagnostic_candidate_indices(
+    potentials: Any,
+    selected_indices: np.ndarray,
+    *,
+    candidate_oracle_rollouts: int,
+    torch: Any,
+) -> Any:
+    if potentials.ndim != 2:
+        raise ValueError("potentials must have shape [B,K]")
+    batch_size, candidate_count = potentials.shape
+    if candidate_oracle_rollouts <= 0:
+        return torch.empty(
+            batch_size,
+            0,
+            dtype=torch.long,
+            device=potentials.device,
+        )
+    selected = torch.as_tensor(
+        selected_indices,
+        dtype=torch.long,
+        device=potentials.device,
+    ).reshape(-1)
+    if selected.shape[0] != batch_size:
+        raise ValueError("selected_indices must have shape [B]")
+    rows: list[list[int]] = []
+    for row_index in range(batch_size):
+        selected_index = int(selected[row_index].item())
+        if not 0 <= selected_index < candidate_count:
+            selected_index = 0
+        chosen = [selected_index]
+        ordered = torch.argsort(potentials[row_index], descending=True)
+        for value in ordered:
+            candidate_index = int(value.item())
+            if candidate_index in chosen:
+                continue
+            if not bool(torch.isfinite(potentials[row_index, candidate_index])):
+                continue
+            chosen.append(candidate_index)
+            if len(chosen) >= candidate_oracle_rollouts:
+                break
+        while len(chosen) < candidate_oracle_rollouts:
+            chosen.append(selected_index)
+        rows.append(chosen)
+    return torch.tensor(rows, dtype=torch.long, device=potentials.device)
+
+
 def _select_action_chunk(
     model: DoVLAModel,
     observations: Any,
@@ -1691,6 +1990,49 @@ def _select_lattice_action_chunk(
                 device=candidate_type_bonus.device,
             )
             candidate_type_bonus = torch.cat([baseline_bonus, candidate_type_bonus], dim=1)
+    candidates, potentials = _score_lattice_action_chunks(
+        model,
+        observations,
+        instructions,
+        action_candidates,
+        torch=torch,
+        action_low=action_low,
+        action_high=action_high,
+        candidate_mask=candidate_mask,
+        candidate_type_bonus=candidate_type_bonus,
+        action_l2_penalty=action_l2_penalty,
+        action_l2_penalty_base=action_l2_penalty_base,
+    )
+    batch_size, candidate_count = potentials.shape
+    best_index = torch.argmax(potentials, dim=1)
+    if selection_margin > 0.0 and candidate_count > 0:
+        baseline = potentials[:, 0]
+        best_potential = potentials.gather(1, best_index.reshape(batch_size, 1)).reshape(
+            batch_size
+        )
+        keep_baseline = best_potential <= baseline + float(selection_margin)
+        best_index = torch.where(keep_baseline, torch.zeros_like(best_index), best_index)
+    batch_index = torch.arange(batch_size, device=candidates.device)
+    best_actions = candidates[batch_index, best_index]
+    return best_actions, best_index.detach().cpu().numpy()
+
+
+def _score_lattice_action_chunks(
+    model: DoVLAModel,
+    observations: Any,
+    instructions: list[str],
+    action_candidates: Any,
+    *,
+    torch: Any,
+    action_low: Any | None,
+    action_high: Any | None,
+    candidate_mask: Any | None,
+    candidate_type_bonus: Any | None = None,
+    action_l2_penalty: float = 0.0,
+    action_l2_penalty_base: Any | None = None,
+) -> tuple[Any, Any]:
+    if action_candidates.ndim != 4:
+        raise ValueError("action_candidates must have shape [B,K,H,D]")
     batch_size, candidate_count = action_candidates.shape[:2]
     candidates = _clamp_action_tensor(
         action_candidates,
@@ -1729,18 +2071,10 @@ def _select_lattice_action_chunk(
         penalty = delta.reshape(batch_size, candidate_count, -1).pow(2).mean(dim=2)
         potentials = potentials - float(action_l2_penalty) * penalty
     if candidate_mask is not None:
+        if candidate_mask.shape != potentials.shape:
+            raise ValueError("candidate_mask must have shape [B,K]")
         potentials = potentials.masked_fill(~candidate_mask, float("-inf"))
-    best_index = torch.argmax(potentials, dim=1)
-    if selection_margin > 0.0 and candidate_count > 0:
-        baseline = potentials[:, 0]
-        best_potential = potentials.gather(1, best_index.reshape(batch_size, 1)).reshape(
-            batch_size
-        )
-        keep_baseline = best_potential <= baseline + float(selection_margin)
-        best_index = torch.where(keep_baseline, torch.zeros_like(best_index), best_index)
-    batch_index = torch.arange(batch_size, device=candidates.device)
-    best_actions = candidates[batch_index, best_index]
-    return best_actions, best_index.detach().cpu().numpy()
+    return candidates, potentials
 
 
 def _select_residual_lattice_action_chunk(
@@ -1792,6 +2126,57 @@ def _select_residual_lattice_action_chunk(
             selection_margin=selection_margin,
             action_l2_penalty=action_l2_penalty,
         )
+    candidates, candidate_mask, candidate_type_bonus = _build_residual_lattice_candidates(
+        policy_mean,
+        residuals,
+        torch=torch,
+        residual_scale=residual_scale,
+        residual_scales=residual_scales,
+        candidate_mask=candidate_mask,
+        candidate_type_bonus=candidate_type_bonus,
+        num_gaussian_candidates=num_gaussian_candidates,
+        candidate_sigma=candidate_sigma,
+        selection_seed=selection_seed,
+    )
+    return _select_lattice_action_chunk(
+        model,
+        observations,
+        instructions,
+        candidates,
+        torch=torch,
+        action_low=(
+            action_low.unsqueeze(1)
+            if action_low is not None and action_low.ndim == 3
+            else action_low
+        ),
+        action_high=(
+            action_high.unsqueeze(1)
+            if action_high is not None and action_high.ndim == 3
+            else action_high
+        ),
+        candidate_mask=candidate_mask,
+        candidate_type_bonus=candidate_type_bonus,
+        selection_margin=selection_margin,
+        action_l2_penalty=action_l2_penalty,
+        action_l2_penalty_base=policy_mean,
+    )
+
+
+def _build_residual_lattice_candidates(
+    policy_mean: Any,
+    residuals: Any,
+    *,
+    torch: Any,
+    residual_scale: float,
+    residual_scales: tuple[float, ...],
+    candidate_mask: Any | None,
+    candidate_type_bonus: Any | None,
+    num_gaussian_candidates: int,
+    candidate_sigma: float,
+    selection_seed: int,
+) -> tuple[Any, Any | None, Any | None]:
+    residuals = _align_residual_horizon_to_policy(policy_mean, residuals, torch=torch)
+    scales = tuple(float(scale) for scale in residual_scales) or (float(residual_scale),)
     candidate_blocks = [
         policy_mean.unsqueeze(1) + scale * residuals
         for scale in scales
@@ -1833,28 +2218,32 @@ def _select_residual_lattice_action_chunk(
                 device=candidate_type_bonus.device,
             )
             candidate_type_bonus = torch.cat([candidate_type_bonus, extra_bonus], dim=1)
-    return _select_lattice_action_chunk(
-        model,
-        observations,
-        instructions,
-        candidates,
-        torch=torch,
-        action_low=(
-            action_low.unsqueeze(1)
-            if action_low is not None and action_low.ndim == 3
-            else action_low
-        ),
-        action_high=(
-            action_high.unsqueeze(1)
-            if action_high is not None and action_high.ndim == 3
-            else action_high
-        ),
-        candidate_mask=candidate_mask,
-        candidate_type_bonus=candidate_type_bonus,
-        selection_margin=selection_margin,
-        action_l2_penalty=action_l2_penalty,
-        action_l2_penalty_base=policy_mean,
+    return candidates, candidate_mask, candidate_type_bonus
+
+
+def _align_residual_horizon_to_policy(policy_mean: Any, residuals: Any, *, torch: Any) -> Any:
+    if residuals.ndim != 4:
+        raise ValueError("residuals must have shape [B,K,H,D]")
+    if policy_mean.ndim != 3:
+        raise ValueError("policy_mean must have shape [B,H,D]")
+    policy_horizon = int(policy_mean.shape[1])
+    residual_horizon = int(residuals.shape[2])
+    if residual_horizon == policy_horizon:
+        return residuals
+    if residual_horizon > policy_horizon:
+        return residuals[:, :, :policy_horizon, :]
+    pad_shape = (
+        residuals.shape[0],
+        residuals.shape[1],
+        policy_horizon - residual_horizon,
+        residuals.shape[3],
     )
+    padding = torch.zeros(
+        pad_shape,
+        dtype=residuals.dtype,
+        device=residuals.device,
+    )
+    return torch.cat([residuals, padding], dim=2)
 
 
 def _select_field_softmax_residual_action_chunk(
@@ -2316,7 +2705,7 @@ def _clip_to_action_space(actions: np.ndarray, env: Any) -> np.ndarray:
 
 
 def _summarize_rows(rows: list[dict[str, Any]]) -> dict[str, Any]:
-    return {
+    summary = {
         "num_groups": len(rows),
         "policy_rollout_success_rate": _mean([row["success"] for row in rows]),
         "policy_rollout_progress": _mean([row["progress"] for row in rows]),
@@ -2331,6 +2720,48 @@ def _summarize_rows(rows: list[dict[str, Any]]) -> dict[str, Any]:
         "action_mse_to_best": _mean([row["action_mse_to_best"] for row in rows]),
         "restore_max_error": max([row["restore_error"] for row in rows], default=0.0),
     }
+    candidate_oracle_rows = [
+        row for row in rows if row.get("candidate_oracle_success") is not None
+    ]
+    if candidate_oracle_rows:
+        summary.update(_candidate_oracle_summary(candidate_oracle_rows))
+    return summary
+
+
+def _candidate_oracle_summary(rows: list[dict[str, Any]]) -> dict[str, Any]:
+    return {
+        "candidate_oracle_success_rate": _mean(
+            [row["candidate_oracle_success"] for row in rows]
+        ),
+        "candidate_oracle_progress": _mean(
+            [row["candidate_oracle_progress"] for row in rows]
+        ),
+        "candidate_oracle_score": _mean(
+            [row["candidate_oracle_score"] for row in rows]
+        ),
+        "candidate_oracle_regret": _mean(
+            [row["candidate_oracle_regret"] for row in rows]
+        ),
+        "candidate_oracle_score_gain_over_selected": _mean(
+            [row["candidate_oracle_score_gain_over_selected"] for row in rows]
+        ),
+        "candidate_oracle_improvement_rate": _mean(
+            [row["candidate_oracle_improves_selected"] for row in rows]
+        ),
+        "candidate_oracle_selected_branch_success_rate": _mean(
+            [row["candidate_oracle_selected_branch_success"] for row in rows]
+        ),
+        "candidate_oracle_selected_branch_progress": _mean(
+            [row["candidate_oracle_selected_branch_progress"] for row in rows]
+        ),
+        "candidate_oracle_restore_max_error": max(
+            [row["candidate_oracle_restore_error"] for row in rows],
+            default=0.0,
+        ),
+        "candidate_oracle_type_counts": dict(
+            Counter(row["candidate_oracle_candidate_type"] for row in rows)
+        ),
+    }
 
 
 def _mean(values: list[Any]) -> float: