Auto-sync: 2026-06-30 09:33:44

dovla_cil/eval/maniskill_policy_rollout.py

@@ -144,6 +144,12 @@ def evaluate_maniskill_policy_rollout(
     field, and executes the best chunk. This keeps the proposal geometry counterfactual while
     avoiding same-state validation candidates.
 
+    When ``selection_mode == 'proposal_lattice'`` the model generates a fixed typed
+    counterfactual proposal set directly from the current state and instruction. The learned
+    field scores those model-generated proposals and executes one selected chunk, preserving
+    clean deployment while testing whether proposal support, not field scoring, is the
+    bottleneck.
+
     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
@@ -173,13 +179,15 @@ def evaluate_maniskill_policy_rollout(
         "policy",
         "field",
         "field_optim",
+        "proposal_lattice",
         "lattice",
         "retrieval_lattice",
         "retrieval_residual",
     }:
         raise ValueError(
-            "selection_mode must be 'policy', 'field', 'field_optim', 'lattice', "
-            "'retrieval_lattice', or 'retrieval_residual'"
+            "selection_mode must be 'policy', 'field', 'field_optim', "
+            "'proposal_lattice', 'lattice', 'retrieval_lattice', or "
+            "'retrieval_residual'"
         )
     if num_candidates <= 0:
         raise ValueError("num_candidates must be positive")
@@ -388,7 +396,12 @@ def evaluate_maniskill_policy_rollout(
             task_summaries[task_id] = _summarize_rows(task_rows)
 
     effective_num_candidates = num_candidates
-    if selection_mode in {"lattice", "retrieval_lattice", "retrieval_residual"}:
+    if selection_mode in {
+        "proposal_lattice",
+        "lattice",
+        "retrieval_lattice",
+        "retrieval_residual",
+    }:
         effective_num_candidates = max(
             [int(row.get("lattice_candidate_count", 0)) for row in rows],
             default=0,
@@ -410,6 +423,9 @@ def evaluate_maniskill_policy_rollout(
         if selection_mode in {"field", "field_optim", "retrieval_residual"}
         and num_candidates > 1
         else 0.0,
+        "proposal_types": list(model_config.proposal_types)
+        if selection_mode == "proposal_lattice"
+        else [],
         "selection_margin": selection_margin,
         "prepend_policy_candidate": bool(prepend_policy_candidate),
         "field_optim_steps": field_optim_steps
@@ -1413,6 +1429,7 @@ def _evaluate_task_cases(
                                 retrieval_residual_scales
                             ),
                             retrieval_residual_reduce=retrieval_residual_reduce,
+                            proposal_types=model_config.proposal_types,
                         ),
                         "selected_residual_scale": _selected_residual_scale(
                             case,
@@ -1433,6 +1450,7 @@ def _evaluate_task_cases(
                                 retrieval_residual_scales
                             ),
                             retrieval_residual_reduce=retrieval_residual_reduce,
+                            proposal_type_count=len(model_config.proposal_types),
                         ),
                         "candidate_source_group_id": case.candidate_source_group_id,
                     }
@@ -1848,6 +1866,29 @@ def _select_action_chunk(
 
     policy_mean = model.forward_policy(observations, instructions)
     batch_size = policy_mean.shape[0]
+    if selection_mode == "proposal_lattice":
+        proposals = model.forward_proposals(observations, instructions)
+        return _select_lattice_action_chunk(
+            model,
+            observations,
+            instructions,
+            proposals,
+            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=None,
+            candidate_type_bonus=None,
+            selection_margin=selection_margin,
+            baseline_action=policy_mean if prepend_policy_candidate else None,
+        )
     if selection_mode == "retrieval_residual":
         if action_candidates is None:
             raise ValueError("retrieval_residual selection requires action_candidates")
@@ -2639,7 +2680,10 @@ def _effective_lattice_candidate_count(
     prepended_policy_candidate: bool = False,
     residual_scale_count: int = 1,
     retrieval_residual_reduce: str = "none",
+    proposal_type_count: int = 0,
 ) -> int:
+    if selection_mode == "proposal_lattice":
+        return int(proposal_type_count) + (1 if prepended_policy_candidate else 0)
     count = len(case.candidate_action_values)
     if prepended_policy_candidate and selection_mode in {"lattice", "retrieval_lattice"}:
         count += 1
@@ -2804,6 +2848,7 @@ def _selected_candidate_type(
     prepended_policy_candidate: bool = False,
     residual_scale_count: int = 1,
     retrieval_residual_reduce: str = "none",
+    proposal_types: tuple[str, ...] = (),
 ) -> str:
     if selection_mode == "policy":
         return "policy_continuous"
@@ -2811,6 +2856,14 @@ def _selected_candidate_type(
         return "field_selected"
     if selection_mode == "field_optim":
         return "field_optim_selected"
+    if selection_mode == "proposal_lattice":
+        if prepended_policy_candidate:
+            if selected_index == 0:
+                return "policy_continuous"
+            selected_index -= 1
+        if 0 <= selected_index < len(proposal_types):
+            return f"proposal_{proposal_types[selected_index]}"
+        return "proposal_unknown"
     if selection_mode == "retrieval_residual":
         if retrieval_residual_reduce in _FIELD_CONDITIONED_RESIDUAL_REDUCERS:
             scale_count = max(1, int(residual_scale_count))

dovla_cil/models/dovla.py

@@ -42,6 +42,7 @@ class DoVLAConfig:
     backbone_model: str | None = None
     backbone_freeze: bool = True
     backbone_local_files_only: bool = True
+    proposal_types: tuple[str, ...] = ()
     # Backward-compatible aliases from the first scaffold.
     observation_dim: int | None = None
     language_dim: int | None = None
@@ -51,6 +52,14 @@ class DoVLAConfig:
             self.obs_dim = int(self.observation_dim)
         if self.language_dim is not None:
             self.lang_dim = int(self.language_dim)
+        if isinstance(self.proposal_types, str):
+            self.proposal_types = tuple(
+                item.strip()
+                for item in self.proposal_types.split(",")
+                if item.strip()
+            )
+        else:
+            self.proposal_types = tuple(str(item) for item in self.proposal_types)
         for name in (
             "obs_dim",
             "lang_dim",
@@ -247,6 +256,25 @@ if nn is not None:
                 self.config.action_dim,
                 action_horizon=self.config.action_horizon,
             )
+            self.proposal_types = tuple(self.config.proposal_types)
+            if self.proposal_types:
+                self.proposal_type_to_index = {
+                    proposal_type: index
+                    for index, proposal_type in enumerate(self.proposal_types)
+                }
+                self.proposal_type_embedding = nn.Embedding(
+                    len(self.proposal_types),
+                    self.config.hidden_dim,
+                )
+                self.proposal_head = PolicyHead(
+                    self.config.hidden_dim,
+                    self.config.action_dim,
+                    action_horizon=self.config.action_horizon,
+                )
+            else:
+                self.proposal_type_to_index = {}
+                self.proposal_type_embedding = None
+                self.proposal_head = None
             self.effect_head = EffectHead(self.config.intervention_dim, self.config.effect_dim)
             self.reward_head = RewardHead(self.config.intervention_dim)
             self.regret_head = RegretHead(self.config.intervention_dim)
@@ -280,6 +308,41 @@ if nn is not None:
             context = self.encode_context(observation, instruction)
             return self.policy_head(context)
 
+        def forward_proposals(self, observation, instruction, proposal_types=None):
+            """Generate a typed counterfactual proposal lattice for each state.
+
+            The standard policy head stays unimodal. This optional head predicts one action
+            chunk per configured intervention family, giving the learned field a clean
+            deployment-time proposal set without reading same-state dataset candidates.
+            """
+
+            if not self.proposal_types or self.proposal_type_embedding is None:
+                raise ValueError("DoVLAConfig.proposal_types is empty")
+            requested = tuple(proposal_types) if proposal_types is not None else self.proposal_types
+            missing = [
+                proposal_type
+                for proposal_type in requested
+                if proposal_type not in self.proposal_type_to_index
+            ]
+            if missing:
+                raise ValueError(f"unknown proposal_types: {missing}")
+            context = self.encode_context(observation, instruction)
+            indices = torch.tensor(
+                [self.proposal_type_to_index[proposal_type] for proposal_type in requested],
+                dtype=torch.long,
+                device=context.device,
+            )
+            type_features = self.proposal_type_embedding(indices)
+            fused = context.unsqueeze(1) + type_features.unsqueeze(0)
+            flat = fused.reshape(context.shape[0] * len(requested), context.shape[-1])
+            proposals = self.proposal_head(flat)
+            return proposals.reshape(
+                context.shape[0],
+                len(requested),
+                self.config.action_horizon,
+                self.config.action_dim,
+            )
+
         def forward_effect(self, observation, instruction, action, skill_type=None):
             z = self.encode_intervention(observation, instruction, action, skill_type=skill_type)
             return self.interventional_field(z)

dovla_cil/training/losses.py

@@ -373,6 +373,7 @@ class InterventionalLossWeights:
     field_preference: float = 1.0
     field_effect: float = 1.0
     field_anchor: float = 0.25
+    proposal: float = 0.0
     # Backward-compatible aliases from the first config schema.
     bc_best_action: float = 1.0
     forward_effect_prediction: float = 1.0

dovla_cil/training/trainer.py

@@ -67,6 +67,7 @@ class TrainerConfig:
     backbone_feature_batch_size: int = 64
     policy_target_types: tuple[str, ...] = ()
     policy_target_map: str | Path | None = None
+    proposal_types: tuple[str, ...] = ()
 
     def __post_init__(self) -> None:
         if self.lr is not None:
@@ -106,6 +107,14 @@ class TrainerConfig:
                 for item in self.policy_target_types.split(",")
                 if item.strip()
             )
+        if isinstance(self.proposal_types, str):
+            self.proposal_types = tuple(
+                item.strip()
+                for item in self.proposal_types.split(",")
+                if item.strip()
+            )
+        else:
+            self.proposal_types = tuple(str(item) for item in self.proposal_types)
         if self.policy_target_map is not None:
             self.policy_target_map = Path(self.policy_target_map)
 
@@ -141,6 +150,7 @@ class DoVLATrainer:
             backbone_model=config.backbone_model,
             backbone_freeze=config.backbone_freeze,
             backbone_local_files_only=config.backbone_local_files_only,
+            proposal_types=config.proposal_types,
         )
         self.model = None
         self.optimizer = None
@@ -301,6 +311,7 @@ class DoVLATrainer:
         pred_best_actions = self.model.forward_policy(best_obs, best_instructions)
 
         bc = behavior_cloning_loss(pred_best_actions, best_actions)
+        proposal = bc * 0.0
         effect = effect_prediction_loss(effect_outputs["effect_vector"], effect_target)
         success = success_loss(effect_outputs["success_logit"], target_success)
         progress = progress_loss(effect_outputs["progress"], target_progress)
@@ -317,10 +328,13 @@ class DoVLATrainer:
             action_features=action,
             neighbors_per_node=self.config.lattice_neighbors,
         )
+        if self.config.proposal_types and float(weights.proposal) > 0.0:
+            proposal = self._proposal_bc_loss(records)
         if self.config.objective == "lattice_field":
             anchor = float(weights.field_anchor)
             total = (
                 weights.weight("bc") * bc
+                + float(weights.proposal) * proposal
                 + float(weights.field_potential) * field["potential"]
                 + float(weights.field_preference) * field["preference"]
                 + float(weights.field_effect) * field["effect"]
@@ -334,6 +348,7 @@ class DoVLATrainer:
         else:
             total = (
                 weights.weight("bc") * bc
+                + float(weights.proposal) * proposal
                 + weights.weight("effect") * effect
                 + weights.weight("success") * success
                 + weights.weight("progress") * progress
@@ -343,6 +358,7 @@ class DoVLATrainer:
         metrics = self._batch_metrics(
             total=total,
             bc=bc,
+            proposal=proposal,
             rank=rank,
             pred_reward=pred_reward,
             pred_progress=effect_outputs["progress"],
@@ -360,6 +376,62 @@ class DoVLATrainer:
         )
         return total, metrics
 
+    def _proposal_bc_loss(self, records: list[CILRecord]) -> Any:
+        assert torch is not None
+        assert self.model is not None
+        grouped: dict[str, CILRecord] = {}
+        for record in records:
+            grouped.setdefault(record.group_id, record)
+        input_records: list[CILRecord] = []
+        target_values: list[list[list[list[float]]]] = []
+        target_mask: list[list[float]] = []
+        zero_action = [
+            [0.0 for _ in range(self.config.action_dim)]
+            for _ in range(self.config.action_horizon)
+        ]
+        for group_id, fallback_record in grouped.items():
+            try:
+                group_records = self.dataset.get_group(group_id)
+            except KeyError:
+                group_records = [record for record in records if record.group_id == group_id]
+            input_records.append(group_records[0] if group_records else fallback_record)
+            group_targets: list[list[list[float]]] = []
+            group_mask: list[float] = []
+            for proposal_type in self.config.proposal_types:
+                candidates = [
+                    record
+                    for record in group_records
+                    if record.candidate_type == proposal_type
+                ]
+                if candidates:
+                    target = max(candidates, key=_record_score)
+                    group_targets.append(
+                        vectorize_toy_action(
+                            target.action_chunk,
+                            action_dim=self.config.action_dim,
+                            action_horizon=self.config.action_horizon,
+                        )
+                    )
+                    group_mask.append(1.0)
+                else:
+                    group_targets.append(zero_action)
+                    group_mask.append(0.0)
+            target_values.append(group_targets)
+            target_mask.append(group_mask)
+
+        if not input_records or not any(any(row) for row in target_mask):
+            return self.model.forward_policy(
+                self._obs_tensor(records[:1]),
+                [records[0].instruction] if records else [""],
+            ).sum() * 0.0
+        predictions = self.model.forward_proposals(
+            self._obs_tensor(input_records),
+            [record.instruction for record in input_records],
+        )
+        targets = torch.tensor(target_values, dtype=torch.float32, device=self.device)
+        mask = torch.tensor(target_mask, dtype=torch.float32, device=self.device)
+        return behavior_cloning_loss(predictions, targets, mask=mask)
+
     def _policy_bc_targets(self, records: list[CILRecord]) -> tuple[list[CILRecord], Any]:
         assert torch is not None
         target_records: list[CILRecord] = []
@@ -421,6 +493,7 @@ class DoVLATrainer:
         *,
         total,
         bc,
+        proposal,
         rank,
         pred_reward,
         pred_progress,
@@ -446,6 +519,7 @@ class DoVLATrainer:
             return {
                 "total_loss": float(total.detach().cpu()),
                 "bc_loss": float(bc.detach().cpu()),
+                "proposal_loss": float(proposal.detach().cpu()),
                 "rank_loss": float(rank.detach().cpu()),
                 "rank_acc": float(rank_acc),
                 "progress_mae": float(progress_mae),
@@ -686,6 +760,7 @@ class _MetricAccumulator:
             return {
                 "total_loss": 0.0,
                 "bc_loss": 0.0,
+                "proposal_loss": 0.0,
                 "rank_loss": 0.0,
                 "rank_acc": 0.0,
                 "progress_mae": 0.0,

logs/auto_sync_hf.log

@@ -260,3 +260,4 @@ No files have been modified since last commit. Skipping to prevent empty commit.
 No files have been modified since last commit. Skipping to prevent empty commit.
 No files have been modified since last commit. Skipping to prevent empty commit.
 No files have been modified since last commit. Skipping to prevent empty commit.
+No files have been modified since last commit. Skipping to prevent empty commit.