JiajunLI/roboimi
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: add full attnres vision backbone

Browse Source
This commit is contained in:
Logic
2026-04-05 00:07:59 +08:00
parent a78006808a
commit 2033169840
9 changed files with 546 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files

26
tests/test_imf_vla_agent.py
View File

@@ -422,6 +422,32 @@ class IMFVLAAgentTest(unittest.TestCase):
self.assertEqual(agent.noise_pred_net.constructor_kwargs['cond_dim'], agent.per_step_cond_dim)
self.assertEqual(agent.noise_pred_net.constructor_kwargs['backbone_type'], 'attnres_full')
def test_hydra_config_instantiates_resnet_imf_attnres_with_full_attnres_vision_backbone(self):
cfg = _compose_cfg(
overrides=[
'agent=resnet_imf_attnres',
'agent.vision_backbone.vision_backbone_mode=attnres_resnet',
'agent.vision_backbone.pretrained_backbone_weights=null',
'agent.vision_backbone.input_shape=[3,56,56]',
'agent.vision_backbone.freeze_backbone=false',
'agent.vision_backbone.attnres_stem_dim=16',
'agent.vision_backbone.attnres_stage_dims=[16,32,64,128]',
'agent.vision_backbone.attnres_stage_depths=[1,1,1,1]',
'agent.vision_backbone.attnres_stage_heads=[2,4,4,8]',
'agent.vision_backbone.attnres_stage_kv_heads=[1,1,1,1]',
'agent.vision_backbone.attnres_stage_window_sizes=[7,7,7,7]',
'agent.head.n_layer=1',
'agent.head.n_emb=16',
]
)
with _stub_optional_modules(include_imf_head=True):
agent = instantiate(cfg.agent)
self.assertEqual(agent.vision_encoder.output_dim, 64)
self.assertEqual(agent.per_step_cond_dim, 64 * agent.num_cams + agent.obs_dim)
self.assertEqual(agent.noise_pred_net.constructor_kwargs['cond_dim'], agent.per_step_cond_dim)
if __name__ == '__main__':
unittest.main()