Auto-sync: 2026-06-30 09:33:44 (part 4)

tests/test_maniskill_policy_rollout.py

@@ -209,10 +209,11 @@ def test_policy_rollout_requires_numeric_action_values() -> None:
 class _StubModel:
     """Minimal stand-in exposing forward_policy / forward_field for selection tests."""
 
-    def __init__(self, torch_module, mean, best_offset):
+    def __init__(self, torch_module, mean, best_offset, proposals=None):
         self._torch = torch_module
         self._mean = mean
         self._best_offset = best_offset
+        self._proposals = proposals
 
     def forward_policy(self, observation, instruction):
         del observation, instruction
@@ -225,6 +226,12 @@ class _StubModel:
         distance = ((action - target) ** 2).reshape(action.shape[0], -1).sum(dim=1)
         return {"potential": -distance}
 
+    def forward_proposals(self, observation, instruction):
+        del observation, instruction
+        if self._proposals is None:
+            raise AssertionError("stub proposals were not configured")
+        return self._proposals
+
 
 def test_policy_mode_returns_policy_mean() -> None:
     import torch
@@ -288,6 +295,29 @@ def test_field_mode_can_prefer_perturbed_candidate() -> None:
     assert index.tolist()[0] != 0
 
 
+def test_proposal_lattice_mode_scores_model_generated_proposals() -> None:
+    import torch
+
+    mean = torch.zeros(1, 1, 2)
+    proposals = torch.tensor([[[[0.1, 0.0]], [[0.4, 0.4]]]], dtype=torch.float32)
+    offset = torch.tensor([[[0.4, 0.4]]], dtype=torch.float32)
+    model = _StubModel(torch, mean, best_offset=offset, proposals=proposals)
+
+    actions, index = _select_action_chunk(
+        model,
+        observations=torch.zeros(1, 3),
+        instructions=["a"],
+        torch=torch,
+        selection_mode="proposal_lattice",
+        num_candidates=1,
+        candidate_sigma=0.0,
+        selection_seed=0,
+    )
+
+    assert torch.allclose(actions, proposals[:, 1])
+    assert index.tolist() == [1]
+
+
 def test_retrieval_residual_margin_can_abstain_to_policy() -> None:
     import torch
 

tests/test_trainer.py

@@ -6,6 +6,7 @@ from types import SimpleNamespace
 from dovla_cil.data.schema import RewardInfo
 from dovla_cil.generation.pipeline import generate_cil_dataset
 from dovla_cil.tasks.library import built_in_toy_tasks
+from dovla_cil.training.losses import InterventionalLossWeights
 from dovla_cil.training.trainer import (
     DoVLATrainer,
     TrainerConfig,
@@ -56,6 +57,41 @@ def test_trainer_runs_one_epoch_and_writes_checkpoints(tmp_path: Path) -> None:
     assert "best_policy" in metrics
 
 
+def test_trainer_can_supervise_typed_proposal_head(tmp_path: Path) -> None:
+    dataset_dir = tmp_path / "cil"
+    run_dir = tmp_path / "run"
+    generate_cil_dataset(
+        backend="toy",
+        tasks=built_in_toy_tasks()[:2],
+        out_dir=dataset_dir,
+        num_states_per_task=2,
+        k=4,
+        seed=9,
+        shard_size=8,
+        inline_observations=True,
+    )
+
+    result = DoVLATrainer(
+        TrainerConfig(
+            dataset_dir=dataset_dir,
+            output_dir=run_dir,
+            epochs=1,
+            batch_groups=2,
+            records_per_group=4,
+            hidden_dim=64,
+            learning_rate=1e-3,
+            seed=9,
+            device="cpu",
+            proposal_types=("expert", "near_miss"),
+            losses=InterventionalLossWeights(proposal=1.0),
+        )
+    ).train()
+
+    assert "proposal_loss" in result["history"][0]["val"]
+    resolved = read_json(run_dir / "resolved_config.json")
+    assert resolved["proposal_types"] == ["expert", "near_miss"]
+
+
 def test_field_utility_includes_terminal_success_bonus() -> None:
     records = [
         SimpleNamespace(