feat(vla): add SmolVLA conditioning and experiment artifacts

tests/test_imf_vla_agent.py

@@ -384,7 +384,7 @@ def _make_images(batch_size, obs_horizon, per_camera_fill):
 
 
 class IMFVLAAgentTest(unittest.TestCase):
-    def _make_agent(self, pred_horizon=3, obs_horizon=2, num_action_steps=2):
+    def _make_agent(self, pred_horizon=3, obs_horizon=2, num_action_steps=2, inference_steps=1):
         agent_cls, agent_module = _load_imf_agent_class()
         head = _RecordingLinearIMFHead()
         agent = agent_cls(
@@ -397,7 +397,7 @@ class IMFVLAAgentTest(unittest.TestCase):
             pred_horizon=pred_horizon,
             obs_horizon=obs_horizon,
             diffusion_steps=10,
-            inference_steps=1,
+            inference_steps=inference_steps,
             num_cams=len(_CAMERA_NAMES),
             camera_names=_CAMERA_NAMES,
             num_action_steps=num_action_steps,
@@ -405,7 +405,7 @@ class IMFVLAAgentTest(unittest.TestCase):
         )
         return agent, head, agent_module
 
-    def test_compute_loss_matches_imf_objective_and_masks_padded_actions(self):
+    def test_compute_loss_uses_pseudo_huber_by_default_and_masks_padded_actions(self):
         agent, head, agent_module = self._make_agent(pred_horizon=3, obs_horizon=2)
         images = _make_images(
             batch_size=1,
@@ -447,7 +447,8 @@ class IMFVLAAgentTest(unittest.TestCase):
         du_dt = scale * v + 2.0
         compound_velocity = u + (t - r).view(1, 1, 1) * du_dt
         target = noise - actions
-        elementwise_loss = (compound_velocity - target) ** 2
+        error = compound_velocity - target
+        elementwise_loss = torch.sqrt(1.0 + error.square()) - 1.0
         mask = (~action_is_pad).unsqueeze(-1).to(elementwise_loss.dtype)
         expected_loss = (elementwise_loss * mask).sum() / (mask.sum() * elementwise_loss.shape[-1])
 
@@ -457,6 +458,60 @@ class IMFVLAAgentTest(unittest.TestCase):
         self.assertTrue(torch.allclose(head.calls[0]['t'], t_sample))
         self.assertTrue(torch.allclose(head.calls[0]['cond'], cond))
 
+    def test_compute_loss_can_use_mse_for_ablation(self):
+        agent_cls, agent_module = _load_imf_agent_class()
+        head = _RecordingLinearIMFHead()
+        agent = agent_cls(
+            vision_backbone=_StubVisionBackbone(),
+            state_encoder=nn.Identity(),
+            action_encoder=nn.Identity(),
+            head=head,
+            action_dim=2,
+            obs_dim=1,
+            pred_horizon=3,
+            obs_horizon=2,
+            diffusion_steps=10,
+            inference_steps=1,
+            num_cams=len(_CAMERA_NAMES),
+            camera_names=_CAMERA_NAMES,
+            num_action_steps=2,
+            head_type='transformer',
+            loss_type='mse',
+        )
+        images = _make_images(
+            batch_size=1,
+            obs_horizon=2,
+            per_camera_fill={'r_vis': 1.0, 'top': 2.0, 'front': 3.0},
+        )
+        qpos = torch.tensor([[[0.25], [0.75]]], dtype=torch.float32)
+        actions = torch.tensor(
+            [[[1.0, -1.0], [0.5, 0.25], [-0.5, 1.5]]],
+            dtype=torch.float32,
+        )
+        noise = torch.tensor(
+            [[[0.2, -0.4], [0.1, 0.3], [0.5, -0.2]]],
+            dtype=torch.float32,
+        )
+        t_sample = torch.tensor([0.8], dtype=torch.float32)
+        r_sample = torch.tensor([0.25], dtype=torch.float32)
+
+        with mock.patch.object(agent_module.torch, 'randn_like', return_value=noise), \
+             mock.patch.object(agent_module.torch, 'rand', side_effect=[t_sample, r_sample]):
+            loss = agent.compute_loss({'images': images, 'qpos': qpos, 'action': actions})
+
+        cond = torch.tensor([[[1.0, 2.0, 3.0, 0.25], [1.0, 2.0, 3.0, 0.75]]], dtype=torch.float32)
+        cond_term = cond.mean(dim=(1, 2), keepdim=True)
+        z_t = (1 - t_sample.view(1, 1, 1)) * actions + t_sample.view(1, 1, 1) * noise
+        scale = head.scale.detach()
+        u = scale * z_t + r_sample.view(1, 1, 1) + 2.0 * t_sample.view(1, 1, 1) + cond_term
+        v = scale * z_t + 3.0 * t_sample.view(1, 1, 1) + cond_term
+        du_dt = scale * v + 2.0
+        compound_velocity = u + (t_sample - r_sample).view(1, 1, 1) * du_dt
+        target = noise - actions
+        expected_loss = (compound_velocity - target).square().mean()
+
+        self.assertAlmostEqual(loss.item(), expected_loss.item(), places=6)
+
     def test_predict_action_uses_one_step_imf_sampling_and_image_conditioning(self):
         agent, head, agent_module = self._make_agent(pred_horizon=3, obs_horizon=2)
         agent.infer_scheduler = _ForbiddenScheduler()
@@ -501,6 +556,41 @@ class IMFVLAAgentTest(unittest.TestCase):
         self.assertTrue(torch.allclose(head.calls[0]['t'], torch.ones(2)))
         self.assertTrue(torch.allclose(head.calls[0]['cond'], expected_cond))
 
+    def test_predict_action_supports_two_half_step_imf_sampling(self):
+        agent, head, agent_module = self._make_agent(pred_horizon=3, obs_horizon=2, inference_steps=2)
+        agent.infer_scheduler = _ForbiddenScheduler()
+
+        images = _make_images(
+            batch_size=1,
+            obs_horizon=2,
+            per_camera_fill={'r_vis': 10.0, 'top': 20.0, 'front': 30.0},
+        )
+        qpos = torch.tensor([[[1.0], [2.0]]], dtype=torch.float32)
+        initial_noise = torch.tensor([[[1.0, -1.0], [0.0, 2.0], [3.0, -2.0]]], dtype=torch.float32)
+
+        with mock.patch.object(agent_module.torch, 'randn', return_value=initial_noise):
+            predicted_actions = agent.predict_action(images, qpos)
+
+        expected_cond = torch.tensor(
+            [[[10.0, 20.0, 30.0, 1.0], [10.0, 20.0, 30.0, 2.0]]],
+            dtype=torch.float32,
+        )
+        cond_term = expected_cond.mean(dim=(1, 2), keepdim=True)
+        first_step = 0.75 * initial_noise - 1.25 - 0.5 * cond_term
+        expected_actions = 0.75 * first_step - 0.5 - 0.5 * cond_term
+
+        self.assertEqual(predicted_actions.shape, (1, 3, 2))
+        self.assertTrue(torch.allclose(predicted_actions, expected_actions))
+        self.assertEqual(len(head.calls), 2)
+        self.assertTrue(torch.allclose(head.calls[0]['r'], torch.full((1,), 0.5)))
+        self.assertTrue(torch.allclose(head.calls[0]['t'], torch.ones(1)))
+        self.assertTrue(torch.allclose(head.calls[0]['sample'], initial_noise))
+        self.assertTrue(torch.allclose(head.calls[0]['cond'], expected_cond))
+        self.assertTrue(torch.allclose(head.calls[1]['r'], torch.zeros(1)))
+        self.assertTrue(torch.allclose(head.calls[1]['t'], torch.full((1,), 0.5)))
+        self.assertTrue(torch.allclose(head.calls[1]['sample'], first_step))
+        self.assertTrue(torch.allclose(head.calls[1]['cond'], expected_cond))
+
     def test_select_action_only_regenerates_when_action_queue_is_empty(self):
         agent, _head, _agent_module = self._make_agent(pred_horizon=4, obs_horizon=2, num_action_steps=2)
         observation = {
@@ -723,6 +813,8 @@ class IMFVLAAgentTest(unittest.TestCase):
         self.assertTrue(cfg.agent.head.time_as_cond)
         self.assertFalse(cfg.agent.head.causal_attn)
         self.assertEqual(cfg.agent.inference_steps, 1)
+        self.assertEqual(cfg.agent.loss_type, 'pseudo_huber')
+        self.assertEqual(cfg.agent.pseudo_huber_delta, 1.0)
         self.assertEqual(list(cfg.agent.camera_names), list(_CAMERA_NAMES))
 
         with _stub_optional_modules(include_imf_head=True):