fix: support dual decoder relaunch and width-adaptive lewm tokens

2026-04-17 20:11:54 +08:00
parent d8066823e2
commit 4cd33258d2
3 changed files with 16 additions and 5 deletions
--- a/tests/test_imf_vla_agent.py
+++ b/tests/test_imf_vla_agent.py
@@ -1267,6 +1267,10 @@ class IMFVLAAgentTest(unittest.TestCase):
        self.assertEqual(agent.condition_sequence_length, agent.obs_horizon)
        self.assertEqual(agent.noise_pred_net.constructor_kwargs['n_obs_steps'], agent.obs_horizon)
        self.assertEqual(agent.future_decoder.constructor_kwargs['cond_dim'], 288)
+        self.assertEqual(
+            agent.future_decoder.constructor_kwargs['n_obs_steps'],
+            agent.lewm_history_horizon,
+        )
        self.assertEqual(agent.future_query_tokens.shape, (1, 1, 288))