Auto-sync: 2026-06-27 15:04:36

dovla_cil/eval/maniskill_policy_rollout.py

@@ -55,6 +55,11 @@ def evaluate_maniskill_policy_rollout(
     num_candidates: int = 1,
     candidate_sigma: float = 0.2,
     selection_seed: int = 0,
+    field_optim_steps: int = 0,
+    field_optim_step_size: float = 0.05,
+    field_optim_trust_radius: float = 0.5,
+    field_optim_l2_penalty: float = 0.0,
+    retrieval_neighbors: int = 1,
     lattice_exclude_types: tuple[str, ...] = (),
 ) -> dict[str, Any]:
     """Execute a checkpoint policy from restored ManiSkill CIL states.
@@ -80,6 +85,17 @@ def evaluate_maniskill_policy_rollout(
     training-split state with the same task rather than the evaluated state's own lattice. This
     tests whether the field can use reusable intervention proposals without same-state proposal
     leakage.
+
+    When ``selection_mode == 'retrieval_residual'`` the evaluator retrieves counterfactual
+    action residuals (candidate minus expert action) from the nearest training-split state(s),
+    adds those residuals around the current policy mean, scores the resulting local proposal
+    lattice with the learned field, and executes the best chunk. This keeps the proposal
+    geometry counterfactual while avoiding same-state validation candidates.
+
+    When ``selection_mode == 'field_optim'`` the evaluator starts from the policy mean plus
+    optional Gaussian multi-start proposals, performs projected gradient ascent on the learned
+    field potential in action space, and executes the best optimized chunk. This is still
+    deployment-clean: no dataset action candidates or rewards are consulted.
     """
 
     try:
@@ -94,12 +110,30 @@ def evaluate_maniskill_policy_rollout(
 
     if group_batch_size <= 0:
         raise ValueError("group_batch_size must be positive")
-    if selection_mode not in {"policy", "field", "lattice", "retrieval_lattice"}:
+    if selection_mode not in {
+        "policy",
+        "field",
+        "field_optim",
+        "lattice",
+        "retrieval_lattice",
+        "retrieval_residual",
+    }:
         raise ValueError(
-            "selection_mode must be 'policy', 'field', 'lattice', or 'retrieval_lattice'"
+            "selection_mode must be 'policy', 'field', 'field_optim', 'lattice', "
+            "'retrieval_lattice', or 'retrieval_residual'"
         )
     if num_candidates <= 0:
         raise ValueError("num_candidates must be positive")
+    if field_optim_steps < 0:
+        raise ValueError("field_optim_steps must be non-negative")
+    if field_optim_step_size < 0:
+        raise ValueError("field_optim_step_size must be non-negative")
+    if field_optim_trust_radius < 0:
+        raise ValueError("field_optim_trust_radius must be non-negative")
+    if field_optim_l2_penalty < 0:
+        raise ValueError("field_optim_l2_penalty must be non-negative")
+    if retrieval_neighbors <= 0:
+        raise ValueError("retrieval_neighbors must be positive")
     if selection_mode == "policy":
         num_candidates = 1
     checkpoint = torch.load(
@@ -112,6 +146,8 @@ def evaluate_maniskill_policy_rollout(
     model = DoVLAModel(model_config).to(resolved_device)
     load_model_state(model, checkpoint)
     model.eval()
+    for parameter in model.parameters():
+        parameter.requires_grad_(False)
 
     trainer_config = checkpoint.get("trainer_config", {})
     dataset = CILDataset(dataset_dir)
@@ -134,16 +170,26 @@ def evaluate_maniskill_policy_rollout(
         group_ids,
         observation_mode=model_config.observation_mode,
     )
-    if selection_mode == "retrieval_lattice":
+    if selection_mode in {"retrieval_lattice", "retrieval_residual"}:
         if all_groups:
-            raise ValueError("retrieval_lattice requires a held-out validation split")
-        cases = _attach_retrieved_lattice_candidates(
-            dataset,
-            cases,
-            heldout_group_ids=split_group_ids,
-            obs_dim=model_config.obs_dim,
-            observation_mode=model_config.observation_mode,
-        )
+            raise ValueError(f"{selection_mode} requires a held-out validation split")
+        if selection_mode == "retrieval_lattice":
+            cases = _attach_retrieved_lattice_candidates(
+                dataset,
+                cases,
+                heldout_group_ids=split_group_ids,
+                obs_dim=model_config.obs_dim,
+                observation_mode=model_config.observation_mode,
+                retrieval_neighbors=retrieval_neighbors,
+            )
+        else:
+            cases = _attach_retrieved_residual_candidates(
+                dataset,
+                cases,
+                heldout_group_ids=split_group_ids,
+                obs_dim=model_config.obs_dim,
+                observation_mode=model_config.observation_mode,
+            )
     by_task: dict[str, list[_RolloutCase]] = defaultdict(list)
     for case in cases:
         by_task[case.task_id].append(case)
@@ -167,13 +213,17 @@ def evaluate_maniskill_policy_rollout(
                 num_candidates=num_candidates,
                 candidate_sigma=candidate_sigma,
                 selection_seed=selection_seed,
+                field_optim_steps=field_optim_steps,
+                field_optim_step_size=field_optim_step_size,
+                field_optim_trust_radius=field_optim_trust_radius,
+                field_optim_l2_penalty=field_optim_l2_penalty,
                 lattice_exclude_types=lattice_exclude_types,
             )
             rows.extend(task_rows)
             task_summaries[task_id] = _summarize_rows(task_rows)
 
     effective_num_candidates = num_candidates
-    if selection_mode in {"lattice", "retrieval_lattice"}:
+    if selection_mode in {"lattice", "retrieval_lattice", "retrieval_residual"}:
         effective_num_candidates = max(
             [int(row.get("lattice_candidate_count", 0)) for row in rows],
             default=0,
@@ -191,7 +241,24 @@ def evaluate_maniskill_policy_rollout(
         "group_batch_size": group_batch_size,
         "selection_mode": selection_mode,
         "num_candidates": effective_num_candidates,
-        "candidate_sigma": candidate_sigma if selection_mode == "field" else 0.0,
+        "candidate_sigma": candidate_sigma
+        if selection_mode in {"field", "field_optim"}
+        else 0.0,
+        "field_optim_steps": field_optim_steps
+        if selection_mode == "field_optim"
+        else 0,
+        "field_optim_step_size": field_optim_step_size
+        if selection_mode == "field_optim"
+        else 0.0,
+        "field_optim_trust_radius": field_optim_trust_radius
+        if selection_mode == "field_optim"
+        else 0.0,
+        "field_optim_l2_penalty": field_optim_l2_penalty
+        if selection_mode == "field_optim"
+        else 0.0,
+        "retrieval_neighbors": retrieval_neighbors
+        if selection_mode in {"retrieval_lattice", "retrieval_residual"}
+        else 0,
         "lattice_exclude_types": list(lattice_exclude_types),
         "policy_rollout_success_rate": _mean([row["success"] for row in rows]),
         "policy_rollout_progress": _mean([row["progress"] for row in rows]),
@@ -332,6 +399,86 @@ def _attach_retrieved_lattice_candidates(
     return output
 
 
+def _attach_retrieved_residual_candidates(
+    dataset: CILDataset,
+    cases: list[_RolloutCase],
+    *,
+    heldout_group_ids: list[str],
+    obs_dim: int,
+    observation_mode: str,
+    retrieval_neighbors: int,
+) -> list[_RolloutCase]:
+    if observation_mode != "state":
+        raise ValueError("retrieval_residual currently supports state observations only")
+    heldout = set(heldout_group_ids)
+    bank: dict[str, list[tuple[str, np.ndarray, list[list[list[float]]], list[str]]]] = (
+        defaultdict(list)
+    )
+    for group_id in dataset.group_ids:
+        if group_id in heldout:
+            continue
+        records = dataset.get_group(group_id)
+        if not records:
+            continue
+        task_ids = {record.task_id for record in records}
+        if len(task_ids) != 1:
+            continue
+        anchor = next((record for record in records if record.candidate_type == "expert"), records[0])
+        anchor_action = np.asarray(_numeric_action_values(anchor), dtype=np.float32)
+        residuals: list[list[list[float]]] = [np.zeros_like(anchor_action).tolist()]
+        candidate_types = ["policy_residual"]
+        for record in records:
+            if record.record_id == anchor.record_id:
+                continue
+            residual = np.asarray(_numeric_action_values(record), dtype=np.float32) - anchor_action
+            residuals.append(residual.tolist())
+            candidate_types.append(f"residual_{record.candidate_type}")
+        feature = np.asarray(
+            vectorize_toy_observation(records[0].observation_inline or {}, obs_dim=obs_dim),
+            dtype=np.float32,
+        )
+        bank[next(iter(task_ids))].append((group_id, feature, residuals, candidate_types))
+
+    output: list[_RolloutCase] = []
+    for case in cases:
+        candidates = bank.get(case.task_id, [])
+        if not candidates:
+            zero = np.zeros_like(np.asarray(case.best_action_values, dtype=np.float32)).tolist()
+            output.append(
+                replace(
+                    case,
+                    candidate_action_values=[zero],
+                    candidate_types=["policy_residual"],
+                    candidate_source_group_id=None,
+                )
+            )
+            continue
+        query = np.asarray(
+            vectorize_toy_observation(case.observation, obs_dim=obs_dim),
+            dtype=np.float32,
+        )
+        nearest = sorted(
+            candidates,
+            key=lambda item: float(np.mean((item[1] - query) ** 2)),
+        )[:retrieval_neighbors]
+        source_group_ids: list[str] = []
+        residuals: list[list[list[float]]] = []
+        candidate_types: list[str] = []
+        for source_group_id, _feature, source_residuals, source_candidate_types in nearest:
+            source_group_ids.append(source_group_id)
+            residuals.extend(source_residuals)
+            candidate_types.extend(source_candidate_types)
+        output.append(
+            replace(
+                case,
+                candidate_action_values=residuals,
+                candidate_types=candidate_types,
+                candidate_source_group_id=";".join(source_group_ids),
+            )
+        )
+    return output
+
+
 def _evaluate_task_cases(
     task_id: str,
     cases: list[_RolloutCase],
@@ -348,6 +495,10 @@ def _evaluate_task_cases(
     num_candidates: int = 1,
     candidate_sigma: float = 0.2,
     selection_seed: int = 0,
+    field_optim_steps: int = 0,
+    field_optim_step_size: float = 0.05,
+    field_optim_trust_radius: float = 0.5,
+    field_optim_l2_penalty: float = 0.0,
     lattice_exclude_types: tuple[str, ...] = (),
 ) -> list[dict[str, Any]]:
     rows: list[dict[str, Any]] = []
@@ -406,6 +557,10 @@ def _evaluate_task_cases(
                 num_candidates=num_candidates,
                 candidate_sigma=candidate_sigma,
                 selection_seed=selection_seed + start,
+                field_optim_steps=field_optim_steps,
+                field_optim_step_size=field_optim_step_size,
+                field_optim_trust_radius=field_optim_trust_radius,
+                field_optim_l2_penalty=field_optim_l2_penalty,
                 action_low=action_low,
                 action_high=action_high,
                 action_candidates=(
@@ -416,6 +571,7 @@ def _evaluate_task_cases(
                     )
                     if selection_mode == "lattice"
                     or selection_mode == "retrieval_lattice"
+                    or selection_mode == "retrieval_residual"
                     else None
                 ),
                 candidate_mask=(
@@ -425,7 +581,7 @@ def _evaluate_task_cases(
                         device=device,
                         exclude_types=lattice_exclude_types,
                     )
-                    if selection_mode in {"lattice", "retrieval_lattice"}
+                    if selection_mode in {"lattice", "retrieval_lattice", "retrieval_residual"}
                     and lattice_exclude_types
                     else None
                 ),
@@ -494,6 +650,10 @@ def _select_action_chunk(
     num_candidates: int,
     candidate_sigma: float,
     selection_seed: int,
+    field_optim_steps: int = 0,
+    field_optim_step_size: float = 0.05,
+    field_optim_trust_radius: float = 0.5,
+    field_optim_l2_penalty: float = 0.0,
     action_low: Any | None = None,
     action_high: Any | None = None,
     action_candidates: Any | None = None,
@@ -503,8 +663,11 @@ def _select_action_chunk(
 
     In ``policy`` mode this is just the deterministic policy mean. In ``field`` mode we draw
     ``num_candidates`` proposals (mean + Gaussian noise) and keep, per state, the chunk whose
-    learned interventional-field potential is largest. Only the model's own generations are
-    scored, so no dataset candidate ever leaks into the deployed action.
+    learned interventional-field potential is largest. In ``field_optim`` mode those proposals
+    are optimized by projected action-space gradient ascent before scoring. Only the model's
+    own generations are scored, so no dataset candidate ever leaks into the deployed action.
+    In ``retrieval_residual`` mode, train-split counterfactual residuals are translated around
+    the current policy mean before field scoring.
     """
 
     if selection_mode in {"lattice", "retrieval_lattice"}:
@@ -523,9 +686,59 @@ def _select_action_chunk(
 
     policy_mean = model.forward_policy(observations, instructions)
     batch_size = policy_mean.shape[0]
+    if selection_mode == "retrieval_residual":
+        if action_candidates is None:
+            raise ValueError("retrieval_residual selection requires action_candidates")
+        return _select_residual_lattice_action_chunk(
+            model,
+            observations,
+            instructions,
+            policy_mean,
+            action_candidates,
+            torch=torch,
+            action_low=action_low,
+            action_high=action_high,
+            candidate_mask=candidate_mask,
+        )
+
     if selection_mode == "policy" or num_candidates <= 1:
+        if selection_mode == "field_optim":
+            return _select_field_optim_action_chunk(
+                model,
+                observations,
+                instructions,
+                policy_mean,
+                torch=torch,
+                num_candidates=1,
+                candidate_sigma=candidate_sigma,
+                selection_seed=selection_seed,
+                action_low=action_low,
+                action_high=action_high,
+                optim_steps=field_optim_steps,
+                optim_step_size=field_optim_step_size,
+                optim_trust_radius=field_optim_trust_radius,
+                optim_l2_penalty=field_optim_l2_penalty,
+            )
         return policy_mean, np.zeros(batch_size, dtype=np.int64)
 
+    if selection_mode == "field_optim":
+        return _select_field_optim_action_chunk(
+            model,
+            observations,
+            instructions,
+            policy_mean,
+            torch=torch,
+            num_candidates=num_candidates,
+            candidate_sigma=candidate_sigma,
+            selection_seed=selection_seed,
+            action_low=action_low,
+            action_high=action_high,
+            optim_steps=field_optim_steps,
+            optim_step_size=field_optim_step_size,
+            optim_trust_radius=field_optim_trust_radius,
+            optim_l2_penalty=field_optim_l2_penalty,
+        )
+
     generator = torch.Generator(device=policy_mean.device)
     generator.manual_seed(int(selection_seed))
     candidates = [_clamp_action_tensor(policy_mean, action_low=action_low, action_high=action_high)]
@@ -566,6 +779,155 @@ def _select_action_chunk(
     return best_actions, best_index.detach().cpu().numpy()
 
 
+def _select_field_optim_action_chunk(
+    model: DoVLAModel,
+    observations: Any,
+    instructions: list[str],
+    policy_mean: Any,
+    *,
+    torch: Any,
+    num_candidates: int,
+    candidate_sigma: float,
+    selection_seed: int,
+    action_low: Any | None,
+    action_high: Any | None,
+    optim_steps: int,
+    optim_step_size: float,
+    optim_trust_radius: float,
+    optim_l2_penalty: float,
+) -> tuple[Any, np.ndarray]:
+    batch_size = policy_mean.shape[0]
+    base = _clamp_action_tensor(
+        policy_mean.detach(),
+        action_low=action_low,
+        action_high=action_high,
+    )
+    generator = torch.Generator(device=base.device)
+    generator.manual_seed(int(selection_seed))
+
+    candidates = base.unsqueeze(1).repeat(1, num_candidates, 1, 1)
+    if num_candidates > 1 and candidate_sigma > 0:
+        noise = torch.randn(
+            candidates.shape,
+            generator=generator,
+            device=base.device,
+            dtype=base.dtype,
+        )
+        noise[:, 0].zero_()
+        candidates = candidates + candidate_sigma * noise
+    candidates = _project_action_candidates(
+        candidates,
+        base,
+        action_low=action_low,
+        action_high=action_high,
+        trust_radius=optim_trust_radius,
+    )
+
+    for _ in range(optim_steps):
+        candidates = candidates.detach().requires_grad_(True)
+        with torch.enable_grad():
+            objective = _field_candidate_objective(
+                model,
+                observations,
+                instructions,
+                candidates,
+                base,
+                torch=torch,
+                l2_penalty=optim_l2_penalty,
+            )
+            grad = torch.autograd.grad(objective.sum(), candidates, only_inputs=True)[0]
+        with torch.no_grad():
+            step = _normalized_action_gradient(grad, torch=torch) * optim_step_size
+            candidates = _project_action_candidates(
+                candidates + step,
+                base,
+                action_low=action_low,
+                action_high=action_high,
+                trust_radius=optim_trust_radius,
+            )
+
+    with torch.no_grad():
+        objective = _field_candidate_objective(
+            model,
+            observations,
+            instructions,
+            candidates,
+            base,
+            torch=torch,
+            l2_penalty=optim_l2_penalty,
+        )
+        best_index = torch.argmax(objective, dim=1)
+        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 _field_candidate_objective(
+    model: DoVLAModel,
+    observations: Any,
+    instructions: list[str],
+    candidates: Any,
+    base: Any,
+    *,
+    torch: Any,
+    l2_penalty: float,
+) -> Any:
+    if candidates.ndim != 4:
+        raise ValueError("candidates must have shape [B,K,H,D]")
+    batch_size, candidate_count = candidates.shape[:2]
+    flat_candidates = candidates.reshape(
+        batch_size * candidate_count,
+        candidates.shape[-2],
+        candidates.shape[-1],
+    )
+    flat_observations = observations.repeat_interleave(candidate_count, dim=0)
+    flat_instructions = [
+        instruction
+        for instruction in instructions
+        for _ in range(candidate_count)
+    ]
+    field = model.forward_field(flat_observations, flat_instructions, flat_candidates)
+    potential = field["potential"].reshape(batch_size, candidate_count)
+    if l2_penalty <= 0:
+        return potential
+    delta = candidates - base.unsqueeze(1)
+    penalty = delta.reshape(batch_size, candidate_count, -1).pow(2).mean(dim=2)
+    return potential - float(l2_penalty) * penalty
+
+
+def _normalized_action_gradient(grad: Any, *, torch: Any) -> Any:
+    flat = grad.reshape(*grad.shape[:2], -1)
+    scale = flat.abs().amax(dim=2).clamp_min(1e-6).reshape(*grad.shape[:2], 1, 1)
+    return torch.nan_to_num(grad / scale)
+
+
+def _project_action_candidates(
+    candidates: Any,
+    base: Any,
+    *,
+    action_low: Any | None,
+    action_high: Any | None,
+    trust_radius: float,
+) -> Any:
+    if trust_radius > 0:
+        low = base.unsqueeze(1) - float(trust_radius)
+        high = base.unsqueeze(1) + float(trust_radius)
+        candidates = candidates.clamp(min=low, max=high)
+    return _clamp_action_tensor(
+        candidates,
+        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
+        ),
+    )
+
+
 def _select_lattice_action_chunk(
     model: DoVLAModel,
     observations: Any,
@@ -606,6 +968,44 @@ def _select_lattice_action_chunk(
     return best_actions, best_index.detach().cpu().numpy()
 
 
+def _select_residual_lattice_action_chunk(
+    model: DoVLAModel,
+    observations: Any,
+    instructions: list[str],
+    policy_mean: Any,
+    action_residuals: Any,
+    *,
+    torch: Any,
+    action_low: Any | None,
+    action_high: Any | None,
+    candidate_mask: Any | None,
+) -> tuple[Any, np.ndarray]:
+    if action_residuals.ndim != 4:
+        raise ValueError("action_residuals must have shape [B,K,H,D]")
+    candidates = policy_mean.unsqueeze(1) + action_residuals.to(
+        device=policy_mean.device,
+        dtype=policy_mean.dtype,
+    )
+    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,
+    )
+
+
 def _lattice_candidate_mask(
     batch: list[_RolloutCase],
     *,
@@ -633,6 +1033,12 @@ def _selected_candidate_type(
         return "policy_continuous"
     if selection_mode == "field":
         return "field_selected"
+    if selection_mode == "field_optim":
+        return "field_optim_selected"
+    if selection_mode == "retrieval_residual":
+        if 0 <= selected_index < len(case.candidate_types):
+            return f"retrieval_residual_{case.candidate_types[selected_index]}"
+        return "retrieval_residual_unknown"
     if 0 <= selected_index < len(case.candidate_types):
         return f"lattice_{case.candidate_types[selected_index]}"
     return "lattice_unknown"