feat(vla): add SmolVLA conditioned agent with prefix encoder and train/eval validation

Introduce SmolVLA-style VLM prefix encoder backbone and IMF-AttnRes conditioned agent. Add episode-level train/val split, action MSE validation in training, and headless eval support.
2026-05-23 22:35:47 +08:00
parent acbd7c605a
commit d94eb8f70b
15 changed files with 2069 additions and 27 deletions
@@ -129,6 +129,38 @@ class EvalVLAExecutionTest(unittest.TestCase):
            ["r_vis", "top", "front"],
        )

+    def test_resolve_eval_image_resize_shape_prefers_condition_encoder_override(self):
+        cfg = OmegaConf.create(
+            {
+                "agent": {
+                    "condition_encoder": {
+                        "eval_image_resize_shape": None,
+                    },
+                },
+                "data": {
+                    "image_resize_shape": [224, 224],
+                },
+            }
+        )
+
+        self.assertIsNone(eval_vla._resolve_eval_image_resize_shape(cfg))
+
+    def test_resolve_eval_image_resize_shape_prefers_vision_backbone_override(self):
+        cfg = OmegaConf.create(
+            {
+                "agent": {
+                    "vision_backbone": {
+                        "eval_image_resize_shape": [256, 256],
+                    },
+                },
+                "data": {
+                    "image_resize_shape": [224, 224],
+                },
+            }
+        )
+
+        self.assertEqual(eval_vla._resolve_eval_image_resize_shape(cfg), (256, 256))
+
    def test_build_episode_plans_without_box_poses_keeps_serial_sampling_lazy(self):
        plans = eval_vla._build_episode_plans(num_episodes=3)

@@ -168,6 +200,58 @@ class EvalVLAExecutionTest(unittest.TestCase):
            np.array([[[[1.0]]], [[[1.0]]], [[[1.0]]]], dtype=np.float32),
        )

+    def test_prepare_local_policy_batch_keeps_latest_variable_task_when_present(self):
+        queues = eval_vla._new_local_policy_queues(obs_horizon=2)
+        first_observation = {
+            "qpos": torch.tensor([1.0, 2.0], dtype=torch.float32),
+            "images": {"front": torch.tensor([[[1.0]]], dtype=torch.float32)},
+            "task": "pick the red cube",
+        }
+        second_observation = {
+            "qpos": torch.tensor([3.0, 4.0], dtype=torch.float32),
+            "images": {"front": torch.tensor([[[2.0]]], dtype=torch.float32)},
+            "task": "insert the peg into the socket",
+        }
+
+        eval_vla._populate_local_policy_queues(queues, first_observation)
+        eval_vla._populate_local_policy_queues(queues, second_observation)
+        batch = eval_vla._prepare_local_policy_batch(
+            queues,
+            obs_horizon=2,
+            camera_names=["front"],
+        )
+
+        self.assertEqual(batch["task"], ["insert the peg into the socket"])
+
+    def test_prepare_local_policy_batch_omits_task_for_legacy_observations(self):
+        queues = eval_vla._new_local_policy_queues(obs_horizon=2)
+        observation = {
+            "qpos": torch.tensor([1.0, 2.0], dtype=torch.float32),
+            "images": {"front": torch.tensor([[[1.0]]], dtype=torch.float32)},
+        }
+
+        eval_vla._populate_local_policy_queues(queues, observation)
+        batch = eval_vla._prepare_local_policy_batch(
+            queues,
+            obs_horizon=2,
+            camera_names=["front"],
+        )
+
+        self.assertNotIn("task", batch)
+
+    def test_serialize_deserialize_policy_batch_preserves_task(self):
+        batch = {
+            "qpos": torch.zeros(1, 2, 2, dtype=torch.float32),
+            "images": {"front": torch.zeros(1, 2, 1, 1, 1, dtype=torch.float32)},
+            "task": ["pick the red cube", "insert the peg into the socket"],
+        }
+
+        serialized = eval_vla._serialize_policy_batch(batch)
+        deserialized = eval_vla._deserialize_policy_batch(serialized, device="cpu")
+
+        self.assertEqual(serialized["task"], batch["task"])
+        self.assertEqual(deserialized["task"], batch["task"])
+
    def test_enqueue_predicted_actions_uses_executable_slice(self):
        queues = eval_vla._new_local_policy_queues(obs_horizon=2)
        predicted_actions = torch.tensor(