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

feat(vla): add SmolVLA conditioning and experiment artifacts

Browse Source
This commit is contained in:
Logic
2026-07-31 10:11:04 +08:00
parent acbd7c605a
commit 5ae9f5fa48
175 changed files with 317471 additions and 70 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test_eval_vla_headless.py
+36
View File
@@ -598,6 +598,42 @@ class EvalVLAHeadlessTest(unittest.TestCase):
execute_policy_action.assert_called_once()
self.assertEqual(fake_env.reset_calls, [sampled_task_state])
def test_parallel_socket_peg_payloads_do_not_plan_transfer_box_poses(self):
cfg = OmegaConf.create(
{
"agent": {},
"eval": {
"num_episodes": 3,
"num_workers": 2,
"task_name": "sim_air_insert_socket_peg",
},
}
)
with mock.patch.object(
eval_vla,
"sample_transfer_pose",
side_effect=AssertionError(
"sim_air_insert_socket_peg parallel rollout must not sample transfer box poses"
),
):
payloads, active_workers = eval_vla._build_parallel_worker_payloads(
cfg,
{"output_dir": None},
)
self.assertEqual(active_workers, 2)
planned_episodes = [
episode_plan
for payload in payloads
for episode_plan in payload["episode_plans"]
]
self.assertEqual(
[plan["episode_index"] for plan in planned_episodes],
[0, 1, 2],
)
self.assertTrue(all("box_pos" not in plan for plan in planned_episodes))
if __name__ == "__main__":
unittest.main()
tests/test_imf_transformer1d_external_alignment.py
+67
View File
@@ -191,6 +191,73 @@ class IMFTransformer1DExternalAlignmentTest(unittest.TestCase):
self._optim_group_names(external_model, external_groups),
)
def test_attnres_full_supports_multihead_grouped_query_attention(self):
local_module = _load_local_module()
model = local_module.IMFTransformer1D(
input_dim=4,
output_dim=4,
horizon=6,
n_obs_steps=3,
cond_dim=5,
n_layer=1,
n_head=4,
n_kv_head=2,
n_emb=16,
p_drop_emb=0.0,
p_drop_attn=0.0,
causal_attn=False,
time_as_cond=True,
n_cond_layers=0,
backbone_type='attnres_full',
)
model.eval()
attention = model.attnres_backbone.layers[0].fn
self.assertEqual(attention.n_heads, 4)
self.assertEqual(attention.n_kv_heads, 2)
self.assertEqual(attention.d_head, 4)
sample = torch.randn(2, 6, 4)
r = torch.tensor([0.1, 0.4], dtype=torch.float32)
t = torch.tensor([0.7, 0.9], dtype=torch.float32)
cond = torch.randn(2, 3, 5)
with torch.no_grad():
output = model(sample=sample, r=r, t=t, cond=cond)
self.assertEqual(output.shape, (2, 6, 4))
def test_attnres_full_rejects_invalid_multihead_shapes(self):
local_module = _load_local_module()
with self.assertRaisesRegex(ValueError, 'must be divisible by n_head'):
local_module.IMFTransformer1D(
input_dim=4,
output_dim=4,
horizon=6,
n_head=5,
n_emb=16,
backbone_type='attnres_full',
)
with self.assertRaisesRegex(ValueError, 'n_head=4 must be divisible by n_kv_head=3'):
local_module.IMFTransformer1D(
input_dim=4,
output_dim=4,
horizon=6,
n_head=4,
n_kv_head=3,
n_emb=16,
backbone_type='attnres_full',
)
with self.assertRaisesRegex(ValueError, 'requires an even per-head dimension'):
local_module.IMFTransformer1D(
input_dim=4,
output_dim=4,
horizon=6,
n_head=2,
n_kv_head=1,
n_emb=18,
backbone_type='attnres_full',
)
if __name__ == '__main__':
unittest.main()
tests/test_imf_vla_agent.py
+96 -4
View File
@@ -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):
tests/test_simple_robot_dataset_image_loading.py
+25 -3
View File
@@ -12,10 +12,13 @@ from roboimi.vla.data.simpe_robot_dataset import SimpleRobotDataset
class SimpleRobotDatasetImageLoadingTest(unittest.TestCase):
def _write_episode(self, dataset_dir: Path) -> None:
episode_path = dataset_dir / "episode_0.hdf5"
def _write_episode(self, dataset_dir: Path, episode_idx: int = 0, action_offset: float = 0.0) -> None:
episode_path = dataset_dir / f"episode_{episode_idx}.hdf5"
with h5py.File(episode_path, "w") as root:
root.create_dataset("action", data=np.arange(8, dtype=np.float32).reshape(4, 2))
root.create_dataset(
"action",
data=(np.arange(8, dtype=np.float32).reshape(4, 2) + action_offset),
)
root.create_dataset(
"observations/qpos",
data=np.arange(16, dtype=np.float32).reshape(4, 4),
@@ -79,3 +82,22 @@ class SimpleRobotDatasetImageLoadingTest(unittest.TestCase):
fake_cv2.resize.assert_not_called()
self.assertEqual(tuple(sample["observation.front"].shape), (2, 3, 8, 8))
def test_dataset_can_filter_by_episode_indices_and_expose_available_episode_indices(self):
with tempfile.TemporaryDirectory() as tmpdir:
dataset_dir = Path(tmpdir)
self._write_episode(dataset_dir, episode_idx=3, action_offset=0.0)
self._write_episode(dataset_dir, episode_idx=7, action_offset=100.0)
dataset = SimpleRobotDataset(
dataset_dir,
obs_horizon=2,
pred_horizon=3,
camera_names=["front"],
episode_indices=[7],
)
sample = dataset[0]
self.assertEqual(dataset.available_episode_indices, [7])
self.assertEqual(len(dataset), 4)
self.assertEqual(sample["action"][0, 0].item(), 100.0)
tests/test_smolvla_imf_agent.py
+341
View File
@@ -0,0 +1,341 @@
import contextlib
import importlib
import importlib.machinery
import sys
import types
import unittest
from pathlib import Path
from unittest import mock
import torch
from hydra import compose, initialize_config_dir
from hydra.core.global_hydra import GlobalHydra
from hydra.utils import instantiate
from omegaconf import OmegaConf
from torch import nn
_REPO_ROOT = Path(__file__).resolve().parents[1]
_CONFIG_DIR = str((_REPO_ROOT / 'roboimi/vla/conf').resolve())
_CAMERA_NAMES = ('r_vis', 'top', 'front')
_MISSING = object()
class _RecordingHead(nn.Module):
def __init__(self):
super().__init__()
self.scale = nn.Parameter(torch.tensor(0.5))
self.calls = []
@staticmethod
def _broadcast(value, reference):
while value.ndim < reference.ndim:
value = value.unsqueeze(-1)
return value
def forward(self, sample, r, t, cond=None):
self.calls.append({
'sample': sample.detach().clone(),
'r': r.detach().clone(),
't': t.detach().clone(),
'cond': None if cond is None else cond.detach().clone(),
})
cond_term = 0.0 if cond is None else cond.mean(dim=(1, 2), keepdim=True)
return self.scale * sample + self._broadcast(r, sample) + 2.0 * self._broadcast(t, sample) + cond_term
class _TaskAwareConditionEncoder(nn.Module):
output_dim = 4
condition_sequence_length = 3
tokens_per_step = 3
joint_output_dim = 4
camera_names = _CAMERA_NAMES
num_cameras = 3
def __init__(self, **kwargs):
super().__init__()
self.constructor_kwargs = dict(kwargs)
self.bias = nn.Parameter(torch.tensor(0.0))
self.calls = []
def forward(self, images, state, task=None):
self.calls.append({'task': task, 'state': state.detach().clone(), 'image_keys': tuple(images.keys())})
batch_size = state.shape[0]
state_last = state[:, -1, 0]
if isinstance(task, str):
task_lengths = torch.full((batch_size,), float(len(task)), dtype=state.dtype, device=state.device)
else:
task_lengths = torch.tensor([float(len(item)) for item in task], dtype=state.dtype, device=state.device)
image_marker = images['r_vis'][:, -1].mean(dim=(1, 2, 3))
token0 = torch.stack([state_last, task_lengths, image_marker, torch.ones_like(state_last)], dim=-1)
token1 = token0 + 1.0
token2 = token0 + 2.0
return torch.stack([token0, token1, token2], dim=1) + self.bias
class _BF16TaskAwareConditionEncoder(_TaskAwareConditionEncoder):
def forward(self, images, state, task=None):
return super().forward(images, state, task=task).to(dtype=torch.bfloat16)
class _StubIMFHead(nn.Module):
def __init__(self, input_dim, output_dim, horizon, n_obs_steps, cond_dim, **kwargs):
super().__init__()
self.constructor_kwargs = {
'input_dim': input_dim,
'output_dim': output_dim,
'horizon': horizon,
'n_obs_steps': n_obs_steps,
'cond_dim': cond_dim,
**kwargs,
}
self.proj = nn.Linear(input_dim, output_dim)
self.cond_obs_emb = nn.Linear(cond_dim, max(cond_dim, 1))
def forward(self, sample, r, t, cond=None):
return torch.zeros_like(sample)
def get_optim_groups(self, weight_decay):
return [
{'params': [self.proj.weight], 'weight_decay': weight_decay},
{'params': [self.proj.bias, self.cond_obs_emb.weight, self.cond_obs_emb.bias], 'weight_decay': 0.0},
]
@contextlib.contextmanager
def _stub_optional_modules(include_head=False, include_condition_encoder=False):
previous = {}
def inject(name, module):
if name not in previous:
previous[name] = sys.modules.get(name, _MISSING)
sys.modules[name] = module
diffusers_module = types.ModuleType('diffusers')
schedulers_module = types.ModuleType('diffusers.schedulers')
ddpm_module = types.ModuleType('diffusers.schedulers.scheduling_ddpm')
ddim_module = types.ModuleType('diffusers.schedulers.scheduling_ddim')
class _FakeScheduler:
def __init__(self, num_train_timesteps=100, **kwargs):
self.config = types.SimpleNamespace(num_train_timesteps=num_train_timesteps)
ddpm_module.DDPMScheduler = _FakeScheduler
ddim_module.DDIMScheduler = _FakeScheduler
diffusers_module.DDPMScheduler = _FakeScheduler
diffusers_module.DDIMScheduler = _FakeScheduler
diffusers_module.schedulers = schedulers_module
try:
inject('diffusers', diffusers_module)
inject('diffusers.schedulers', schedulers_module)
inject('diffusers.schedulers.scheduling_ddpm', ddpm_module)
inject('diffusers.schedulers.scheduling_ddim', ddim_module)
if include_head:
import roboimi.vla.models.heads as heads_package
head_module = types.ModuleType('roboimi.vla.models.heads.imf_transformer1d')
head_module.IMFTransformer1D = _StubIMFHead
inject('roboimi.vla.models.heads.imf_transformer1d', head_module)
setattr(heads_package, 'imf_transformer1d', head_module)
if include_condition_encoder:
module = types.ModuleType('tests.fake_smolvla_condition_encoder')
module.TaskAwareConditionEncoder = _TaskAwareConditionEncoder
module.BF16TaskAwareConditionEncoder = _BF16TaskAwareConditionEncoder
inject('tests.fake_smolvla_condition_encoder', module)
yield
finally:
for name, old in reversed(list(previous.items())):
if old is _MISSING:
sys.modules.pop(name, None)
else:
sys.modules[name] = old
def _compose_cfg(overrides=None):
if not OmegaConf.has_resolver('len'):
OmegaConf.register_new_resolver('len', lambda x: len(x))
GlobalHydra.instance().clear()
with initialize_config_dir(version_base=None, config_dir=_CONFIG_DIR):
return compose(config_name='config', overrides=list(overrides or []))
class SmolVLAIMFAgentTest(unittest.TestCase):
def test_compute_loss_uses_pseudo_huber_and_passes_variable_task_to_condition_encoder(self):
from roboimi.vla.agent_smolvla_conditioned import SmolVLAIMFAttnResAgent
condition_encoder = _BF16TaskAwareConditionEncoder()
head = _RecordingHead()
agent = SmolVLAIMFAttnResAgent(
condition_encoder=condition_encoder,
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=3,
camera_names=_CAMERA_NAMES,
num_action_steps=2,
head_type='transformer',
)
images = {
'r_vis': torch.full((2, 2, 1, 2, 2), 1.0),
'top': torch.full((2, 2, 1, 2, 2), 2.0),
'front': torch.full((2, 2, 1, 2, 2), 3.0),
}
qpos = torch.tensor([[[1.0], [2.0]], [[3.0], [4.0]]])
actions = torch.zeros(2, 3, 2)
tasks = ['short', 'longer task']
noise = torch.tensor(
[
[[0.2, -0.4], [0.1, 0.3], [0.5, -0.2]],
[[-0.1, 0.25], [0.4, -0.5], [0.2, 0.6]],
],
dtype=torch.float32,
)
t_sample = torch.full((2,), 0.75, dtype=torch.float32)
r_sample = torch.full((2,), 0.25, dtype=torch.float32)
with mock.patch('roboimi.vla.agent_imf.torch.randn_like', return_value=noise), \
mock.patch(
'roboimi.vla.agent_imf.torch.rand',
side_effect=[t_sample, r_sample],
):
loss = agent.compute_loss({'images': images, 'qpos': qpos, 'action': actions, 'task': tasks})
self.assertTrue(torch.isfinite(loss))
self.assertEqual(condition_encoder.calls[-1]['task'], tasks)
self.assertEqual(head.calls[-1]['cond'].shape, (2, 3, 4))
self.assertTrue(torch.allclose(head.calls[-1]['cond'][:, 0, 1], torch.tensor([5.0, 11.0])))
cond = head.calls[0]['cond'].float()
cond_term = cond.mean(dim=(1, 2), keepdim=True)
z_t = (1 - t_sample.view(2, 1, 1)) * actions + t_sample.view(2, 1, 1) * noise
scale = head.scale.detach()
u = scale * z_t + r_sample.view(2, 1, 1) + 2.0 * t_sample.view(2, 1, 1) + cond_term
v = scale * z_t + 3.0 * t_sample.view(2, 1, 1) + cond_term
du_dt = scale * v + 2.0
compound_velocity = u + (t_sample - r_sample).view(2, 1, 1) * du_dt
expected_loss = (torch.sqrt(1.0 + (compound_velocity - noise).square()) - 1.0).mean()
mse_loss = (compound_velocity - noise).square().mean()
self.assertAlmostEqual(loss.item(), expected_loss.item(), places=4)
self.assertLess(loss.item(), mse_loss.item())
with mock.patch('roboimi.vla.agent_imf.torch.randn', return_value=torch.zeros(2, 3, 2)):
pred = agent.predict_action(images, qpos, task=tasks)
self.assertEqual(pred.shape, (2, 3, 2))
self.assertEqual(condition_encoder.calls[-1]['task'], tasks)
def test_condition_tokens_are_cast_to_action_head_dtype(self):
from roboimi.vla.agent_smolvla_conditioned import SmolVLAIMFAttnResAgent
condition_encoder = _TaskAwareConditionEncoder()
head = _RecordingHead()
agent = SmolVLAIMFAttnResAgent(
condition_encoder=condition_encoder,
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=3,
camera_names=_CAMERA_NAMES,
num_action_steps=2,
head_type='transformer',
)
images = {
'r_vis': torch.full((1, 2, 1, 2, 2), 1.0),
'top': torch.full((1, 2, 1, 2, 2), 2.0),
'front': torch.full((1, 2, 1, 2, 2), 3.0),
}
qpos = torch.tensor([[[1.0], [2.0]]], dtype=torch.float32)
cond = agent._build_cond(images, qpos, task=['pick'])
self.assertEqual(cond.dtype, head.scale.dtype)
def test_unknown_dataset_task_uses_configured_task_description(self):
from roboimi.vla.agent_smolvla_conditioned import SmolVLAIMFAttnResAgent
condition_encoder = _TaskAwareConditionEncoder()
head = _RecordingHead()
agent = SmolVLAIMFAttnResAgent(
condition_encoder=condition_encoder,
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=3,
camera_names=_CAMERA_NAMES,
num_action_steps=2,
head_type='transformer',
task_description='insert the peg into the socket',
)
images = {
'r_vis': torch.full((2, 2, 1, 2, 2), 1.0),
'top': torch.full((2, 2, 1, 2, 2), 2.0),
'front': torch.full((2, 2, 1, 2, 2), 3.0),
}
qpos = torch.tensor([[[1.0], [2.0]], [[3.0], [4.0]]], dtype=torch.float32)
agent._build_cond(images, qpos, task=['unknown', ''])
self.assertEqual(
condition_encoder.calls[-1]['task'],
['insert the peg into the socket', 'insert the peg into the socket'],
)
def test_hydra_config_instantiates_smolvla_imf_attnres_with_condition_encoder_contract(self):
cfg = _compose_cfg(overrides=[
'agent=smolvla_imf_attnres',
'agent.condition_encoder._target_=tests.fake_smolvla_condition_encoder.TaskAwareConditionEncoder',
'agent.condition_dim=4',
'agent.condition_sequence_length=3',
'agent.head.n_layer=1',
'agent.head.n_emb=16',
])
self.assertEqual(cfg.agent._target_, 'roboimi.vla.agent_smolvla_conditioned.SmolVLAIMFAttnResAgent')
self.assertEqual(cfg.agent.head.cond_dim, cfg.agent.condition_dim)
self.assertEqual(cfg.agent.head.n_obs_steps, cfg.agent.condition_sequence_length)
with _stub_optional_modules(include_head=True, include_condition_encoder=True):
agent = instantiate(cfg.agent)
self.assertEqual(agent.per_step_cond_dim, 4)
self.assertEqual(agent.condition_sequence_length, 3)
self.assertIsInstance(agent.noise_pred_net, _StubIMFHead)
self.assertEqual(agent.noise_pred_net.constructor_kwargs['cond_dim'], 4)
self.assertEqual(agent.noise_pred_net.constructor_kwargs['n_obs_steps'], 3)
self.assertEqual(agent.loss_type, 'pseudo_huber')
self.assertEqual(agent.pseudo_huber_delta, 1.0)
def test_hydra_config_exposes_smolvla_pretrained_vlm_defaults(self):
cfg = _compose_cfg(overrides=[
'agent=smolvla_imf_attnres',
])
self.assertEqual(cfg.agent.condition_encoder.model_name, 'HuggingFaceTB/SmolVLM2-500M-Video-Instruct')
self.assertTrue(cfg.agent.condition_encoder.load_vlm_weights)
self.assertEqual(cfg.agent.condition_encoder.num_vlm_layers, 16)
self.assertTrue(cfg.agent.condition_encoder.freeze_vlm)
self.assertTrue(cfg.agent.condition_encoder.freeze_vision_encoder)
self.assertTrue(cfg.agent.condition_encoder.run_text_model)
self.assertEqual(cfg.agent.condition_encoder.max_state_dim, 32)
self.assertIsNone(cfg.agent.condition_encoder.dataset_image_resize_shape)
self.assertIsNone(cfg.agent.condition_encoder.eval_image_resize_shape)
self.assertEqual(cfg.agent.condition_dim, 960)
self.assertEqual(cfg.agent.condition_sequence_length, 241)
self.assertEqual(cfg.agent.head.cond_dim, 960)
self.assertEqual(cfg.agent.head.n_obs_steps, 241)
self.assertEqual(cfg.agent.loss_type, 'pseudo_huber')
self.assertEqual(cfg.agent.pseudo_huber_delta, 1.0)
if __name__ == '__main__':
unittest.main()
tests/test_smolvla_prefix_encoder.py
+382
View File
@@ -0,0 +1,382 @@
import types
import unittest
from unittest import mock
import torch
from torch import nn
class _FakeVisionOutput:
def __init__(self, last_hidden_state):
self.last_hidden_state = last_hidden_state
class _FakeVisionModel(nn.Module):
def __init__(self, hidden_size=4):
super().__init__()
self.dtype = torch.float32
self.scale = nn.Parameter(torch.tensor(1.0))
self.calls = []
self.hidden_size = hidden_size
def forward(self, pixel_values=None, patch_attention_mask=None):
self.calls.append({
'pixel_values': pixel_values.detach().clone(),
'patch_attention_mask': patch_attention_mask,
})
pooled = pixel_values.mean(dim=(2, 3)) * self.scale
tokens = torch.stack([pooled, pooled + 1.0], dim=1)
return _FakeVisionOutput(tokens)
class _FakeConnector(nn.Module):
def __init__(self, in_dim=3, out_dim=4):
super().__init__()
self.proj = nn.Linear(in_dim, out_dim, bias=False)
with torch.no_grad():
self.proj.weight.copy_(
torch.tensor(
[
[1.0, 0.0, 0.0],
[0.0, 1.0, 0.0],
[0.0, 0.0, 1.0],
[1.0, 1.0, 1.0],
]
)
)
def forward(self, x):
return self.proj(x)
class _FakeTextModel(nn.Module):
def __init__(self, vocab_size=32, hidden_size=4, num_layers=6):
super().__init__()
self.embed = nn.Embedding(vocab_size, hidden_size)
self.layers = nn.ModuleList([nn.Linear(hidden_size, hidden_size) for _ in range(num_layers)])
self.norm = nn.Identity()
self.forward_calls = []
with torch.no_grad():
for idx in range(vocab_size):
self.embed.weight[idx].fill_(float(idx))
def get_input_embeddings(self):
return self.embed
def forward(
self,
input_ids=None,
attention_mask=None,
position_ids=None,
past_key_values=None,
inputs_embeds=None,
use_cache=None,
return_dict=True,
cache_position=None,
**kwargs,
):
del input_ids, past_key_values, use_cache, cache_position, kwargs
self.forward_calls.append({
'attention_mask': None if attention_mask is None else attention_mask.detach().clone(),
'position_ids': None if position_ids is None else position_ids.detach().clone(),
'inputs_embeds': inputs_embeds.detach().clone(),
})
hidden = inputs_embeds
for layer in self.layers:
hidden = layer(hidden)
hidden = self.norm(hidden)
if return_dict:
return types.SimpleNamespace(last_hidden_state=hidden)
return (hidden,)
class _FakeVLM(nn.Module):
def __init__(self):
super().__init__()
text_config = types.SimpleNamespace(hidden_size=4, head_dim=2, num_attention_heads=2, num_key_value_heads=1)
self.config = types.SimpleNamespace(text_config=text_config)
self.model = types.SimpleNamespace(
vision_model=_FakeVisionModel(hidden_size=4),
connector=_FakeConnector(in_dim=3, out_dim=4),
text_model=_FakeTextModel(hidden_size=4, num_layers=6),
)
self.lm_head = nn.Linear(4, 32)
class _FakeTokenizer:
fake_image_token_id = 29
global_image_token_id = 30
def __init__(self):
self.padding_side = 'left'
self.calls = []
def __call__(self, text, *, padding, max_length, return_tensors, truncation):
self.calls.append({
'text': list(text),
'padding': padding,
'max_length': max_length,
'return_tensors': return_tensors,
'truncation': truncation,
'padding_side': self.padding_side,
})
batch = len(text)
ids = torch.zeros(batch, max_length, dtype=torch.long)
mask = torch.zeros(batch, max_length, dtype=torch.bool)
for row, item in enumerate(text):
del item
ids[row, :3] = torch.tensor([1, 2, 3])
mask[row, :3] = True
return {'input_ids': ids, 'attention_mask': mask}
class SmolVLAPrefixEncoderTest(unittest.TestCase):
def test_loads_pretrained_vlm_with_local_files_only_crops_layers_and_freezes_vlm(self):
from roboimi.vla.models.backbones.smolvla_prefix_encoder import SmolVLAPrefixEncoder
fake_vlm = _FakeVLM()
fake_tokenizer = _FakeTokenizer()
with mock.patch(
'roboimi.vla.models.backbones.smolvla_prefix_encoder.AutoModelForImageTextToText.from_pretrained',
return_value=fake_vlm,
) as model_loader, mock.patch(
'roboimi.vla.models.backbones.smolvla_prefix_encoder.AutoTokenizer.from_pretrained',
return_value=fake_tokenizer,
) as tokenizer_loader:
encoder = SmolVLAPrefixEncoder(
model_name='HuggingFaceTB/SmolVLM2-500M-Video-Instruct',
local_files_only=True,
num_vlm_layers=2,
freeze_vlm=True,
max_state_dim=5,
tokenizer_max_length=4,
camera_names=('r_vis', 'top'),
resize_imgs_with_padding=None,
)
model_loader.assert_called_once()
self.assertEqual(model_loader.call_args.kwargs['local_files_only'], True)
self.assertEqual(model_loader.call_args.kwargs['torch_dtype'], 'bfloat16')
tokenizer_loader.assert_called_once_with(
'HuggingFaceTB/SmolVLM2-500M-Video-Instruct',
local_files_only=True,
)
self.assertEqual(len(encoder.vlm.model.text_model.layers), 2)
self.assertTrue(all(not p.requires_grad for p in encoder.vlm.parameters()))
self.assertFalse(encoder.vlm.training)
encoder.train()
self.assertFalse(encoder.vlm.training)
def test_accepts_train_eval_resize_compatibility_fields(self):
from roboimi.vla.models.backbones.smolvla_prefix_encoder import SmolVLAPrefixEncoder
encoder = SmolVLAPrefixEncoder(
vlm=_FakeVLM(),
tokenizer=_FakeTokenizer(),
num_vlm_layers=3,
camera_names=('r_vis', 'top'),
resize_imgs_with_padding=None,
dataset_image_resize_shape=None,
eval_image_resize_shape=(640, 480),
)
self.assertIsNone(encoder.dataset_image_resize_shape)
self.assertEqual(encoder.eval_image_resize_shape, (640, 480))
def test_embed_prefix_uses_variable_tasks_camera_order_last_state_and_smolvla_masks(self):
from roboimi.vla.models.backbones.smolvla_prefix_encoder import SmolVLAPrefixEncoder
fake_vlm = _FakeVLM()
fake_tokenizer = _FakeTokenizer()
encoder = SmolVLAPrefixEncoder(
vlm=fake_vlm,
tokenizer=fake_tokenizer,
num_vlm_layers=3,
freeze_vlm=True,
train_state_proj=True,
max_state_dim=5,
tokenizer_max_length=4,
camera_names=('r_vis', 'top'),
resize_imgs_with_padding=None,
)
with torch.no_grad():
encoder.state_proj.weight.zero_()
encoder.state_proj.bias.zero_()
encoder.state_proj.weight[:, :4] = torch.eye(4)
images = {
'top': torch.full((2, 2, 3, 2, 2), 0.75),
'r_vis': torch.full((2, 2, 3, 2, 2), 0.25),
}
state = torch.tensor(
[
[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]],
[[7.0, 8.0, 9.0], [10.0, 11.0, 12.0]],
]
)
tasks = ['pick red cube', 'open drawer']
out = encoder.embed_prefix(images=images, state=state, task=tasks)
# 2 cameras * 2 image tokens + 4 language tokens + 1 state token
self.assertEqual(out.shape, (2, 9, 4))
self.assertEqual(encoder.output_dim, 4)
self.assertEqual(encoder.tokens_per_step, 9)
self.assertEqual(encoder.condition_sequence_length, 9)
self.assertEqual(fake_tokenizer.calls[-1]['text'], ['pick red cube\n', 'open drawer\n'])
self.assertEqual(fake_tokenizer.calls[-1]['padding'], 'max_length')
self.assertEqual(fake_tokenizer.calls[-1]['max_length'], 4)
self.assertEqual(fake_tokenizer.calls[-1]['padding_side'], 'right')
# Camera order is r_vis then top, and pixels are mapped [0,1] -> [-1,1].
first_camera_pixels = fake_vlm.model.vision_model.calls[0]['pixel_values']
second_camera_pixels = fake_vlm.model.vision_model.calls[1]['pixel_values']
self.assertTrue(torch.allclose(first_camera_pixels, torch.full((2, 3, 2, 2), -0.5)))
self.assertTrue(torch.allclose(second_camera_pixels, torch.full((2, 3, 2, 2), 0.5)))
# Last token is padded last state projected from [4,5,6,0,0] and [10,11,12,0,0].
self.assertTrue(torch.allclose(out[0, -1], torch.tensor([4.0, 5.0, 6.0, 0.0])))
self.assertTrue(torch.allclose(out[1, -1], torch.tensor([10.0, 11.0, 12.0, 0.0])))
pad_mask, att_mask = encoder.last_prefix_pad_mask, encoder.last_prefix_att_mask
self.assertEqual(pad_mask.shape, (2, 9))
self.assertEqual(att_mask.shape, (2, 9))
self.assertTrue(torch.all(att_mask[:, :8] == 0))
self.assertTrue(torch.all(att_mask[:, -1] == 1))
def test_forward_can_run_cropped_frozen_text_model_over_prefix_tokens(self):
from roboimi.vla.models.backbones.smolvla_prefix_encoder import SmolVLAPrefixEncoder
fake_vlm = _FakeVLM()
fake_tokenizer = _FakeTokenizer()
encoder = SmolVLAPrefixEncoder(
vlm=fake_vlm,
tokenizer=fake_tokenizer,
num_vlm_layers=2,
freeze_vlm=True,
max_state_dim=5,
tokenizer_max_length=4,
camera_names=('r_vis',),
resize_imgs_with_padding=None,
run_text_model=True,
)
images = {
'r_vis': torch.full((2, 1, 3, 2, 2), 0.25),
}
state = torch.tensor([[[1.0, 2.0, 3.0]], [[4.0, 5.0, 6.0]]])
out = encoder(images, state=state, task=['pick', 'place'])
# 1 camera * 2 image tokens + 4 language tokens + 1 state token.
self.assertEqual(out.shape, (2, 7, 4))
self.assertEqual(len(fake_vlm.model.text_model.layers), 2)
self.assertEqual(len(fake_vlm.model.text_model.forward_calls), 1)
text_call = fake_vlm.model.text_model.forward_calls[-1]
self.assertEqual(tuple(text_call['attention_mask'].shape), (2, 1, 7, 7))
self.assertEqual(tuple(text_call['position_ids'].shape), (2, 7))
self.assertTrue(torch.all(out != 0))
def test_embedding_text_encoder_mode_prunes_text_layers_and_skips_text_forward(self):
from roboimi.vla.models.backbones.smolvla_prefix_encoder import SmolVLAPrefixEncoder
fake_vlm = _FakeVLM()
encoder = SmolVLAPrefixEncoder(
vlm=fake_vlm,
tokenizer=_FakeTokenizer(),
num_vlm_layers=4,
freeze_vlm=True,
max_state_dim=5,
tokenizer_max_length=4,
camera_names=('r_vis',),
resize_imgs_with_padding=None,
text_encoder_mode='embedding',
)
with torch.no_grad():
encoder.state_proj.weight.zero_()
encoder.state_proj.bias.zero_()
encoder.state_proj.weight[:, :4] = torch.eye(4)
images = {
'r_vis': torch.full((2, 1, 3, 2, 2), 0.25),
}
state = torch.tensor([[[1.0, 2.0, 3.0]], [[4.0, 5.0, 6.0]]])
out = encoder(images, state=state, task=['pick', 'place'])
# 1 camera * 2 image tokens + 4 language embedding tokens + 1 state token.
self.assertEqual(out.shape, (2, 7, 4))
self.assertEqual(encoder.text_encoder_mode, 'embedding')
self.assertFalse(encoder.run_text_model)
self.assertEqual(len(fake_vlm.model.text_model.layers), 0)
self.assertEqual(len(fake_vlm.model.text_model.forward_calls), 0)
self.assertIsInstance(fake_vlm.lm_head, nn.Identity)
# Embedding mode returns raw prefix embeddings rather than transformer-scaled inputs.
# For pixel value 0.25 -> prepared -0.5, connector token 0 is [-0.5,-0.5,-0.5,-1.5].
self.assertTrue(torch.allclose(out[0, 0], torch.tensor([-0.5, -0.5, -0.5, -1.5])))
self.assertTrue(torch.allclose(out[0, 2], torch.full((4,), 1.0)))
self.assertTrue(torch.allclose(out[0, -1], torch.tensor([1.0, 2.0, 3.0, 0.0])))
def test_text_encoder_mode_rejects_invalid_value(self):
from roboimi.vla.models.backbones.smolvla_prefix_encoder import SmolVLAPrefixEncoder
with self.assertRaisesRegex(ValueError, 'text_encoder_mode'):
SmolVLAPrefixEncoder(
vlm=_FakeVLM(),
tokenizer=_FakeTokenizer(),
camera_names=('r_vis',),
resize_imgs_with_padding=None,
text_encoder_mode='full',
)
def test_frozen_vlm_still_backpropagates_through_text_model_to_state_projection(self):
from roboimi.vla.models.backbones.smolvla_prefix_encoder import SmolVLAPrefixEncoder
fake_vlm = _FakeVLM()
encoder = SmolVLAPrefixEncoder(
vlm=fake_vlm,
tokenizer=_FakeTokenizer(),
num_vlm_layers=3,
freeze_vlm=True,
train_state_proj=True,
max_state_dim=5,
tokenizer_max_length=4,
camera_names=('r_vis',),
resize_imgs_with_padding=None,
run_text_model=True,
)
images = {
'r_vis': torch.full((2, 1, 3, 2, 2), 0.25),
}
state = torch.tensor([[[1.0, 2.0, 3.0]], [[4.0, 5.0, 6.0]]])
out = encoder(images, state=state, task=['pick', 'place'])
out[:, -1].sum().backward()
self.assertIsNotNone(encoder.state_proj.weight.grad)
self.assertGreater(float(encoder.state_proj.weight.grad.abs().sum()), 0.0)
self.assertTrue(all(param.grad is None for param in fake_vlm.parameters()))
def test_embed_prefix_rejects_missing_camera_and_task_batch_mismatch(self):
from roboimi.vla.models.backbones.smolvla_prefix_encoder import SmolVLAPrefixEncoder
encoder = SmolVLAPrefixEncoder(
vlm=_FakeVLM(),
tokenizer=_FakeTokenizer(),
num_vlm_layers=3,
camera_names=('r_vis', 'top'),
resize_imgs_with_padding=None,
max_state_dim=5,
tokenizer_max_length=4,
)
images = {'r_vis': torch.rand(2, 1, 3, 2, 2)}
state = torch.rand(2, 1, 3)
with self.assertRaisesRegex(ValueError, 'missing.*top'):
encoder.embed_prefix(images=images, state=state, task=['a', 'b'])
images['top'] = torch.rand(2, 1, 3, 2, 2)
with self.assertRaisesRegex(ValueError, 'task batch'):
encoder.embed_prefix(images=images, state=state, task=['only one'])
if __name__ == '__main__':
unittest.main()
tests/test_train_vla_rollout_validation.py
+183
View File
@@ -18,6 +18,17 @@ class _FakeDataset:
return 4
class _FakeEpisodeDataset(_FakeDataset):
def __init__(self, episode_indices=None):
self.available_episode_indices = [100, 101]
self.episode_indices = None if episode_indices is None else list(episode_indices)
def __len__(self):
if self.episode_indices is None:
return 8
return 4 * len(self.episode_indices)
class _FakeLoader:
def __init__(self, batch, length=1):
self._batches = [batch] * length
@@ -90,6 +101,9 @@ class _FakeAgent(nn.Module):
def get_normalization_stats(self):
return {}
def predict_action_chunk(self, agent_input):
return agent_input['action']
class _SequentialLossAgent(nn.Module):
def __init__(self, losses):
@@ -340,6 +354,175 @@ class TrainVLARolloutValidationTest(unittest.TestCase):
self.assertIn('image_resize_shape', captured_dataset_kwargs)
self.assertIsNone(captured_dataset_kwargs['image_resize_shape'])
def test_explicit_val_episode_indices_split_train_and_val_datasets(self):
cfg = OmegaConf.create(
{
'agent': {
'normalization_type': 'min_max',
},
'data': {
'dataset_dir': 'unused',
'camera_names': ['front'],
},
'train': {
'batch_size': 2,
'lr': 1e-4,
'max_steps': 0,
'device': 'cpu',
'disable_cudnn': False,
'num_workers': 0,
'val_split': 0.0,
'val_episode_indices': [100],
'action_mse_val_freq_epochs': 1,
'seed': 42,
'log_freq': 1,
'save_freq': 10,
'use_swanlab': False,
'rollout_val_freq_epochs': 0,
'rollout_validate_on_checkpoint': False,
'rollout_num_episodes': 1,
'warmup_steps': 1,
'scheduler_type': 'constant',
'min_lr': 1e-6,
'weight_decay': 1e-5,
'grad_clip': 1.0,
'pretrained_ckpt': None,
'resume_ckpt': None,
},
'eval': {
'ckpt_path': 'unused.pt',
'num_episodes': 1,
'headless': True,
'device': 'cpu',
'verbose_action': False,
},
'experiment': {},
}
)
instantiate_kwargs = []
def fake_instantiate(config_node, **kwargs):
if config_node is cfg.data:
instantiate_kwargs.append(dict(kwargs))
return _FakeEpisodeDataset(kwargs.get('episode_indices'))
if config_node is cfg.agent:
return _FakeAgent()
raise AssertionError(f'unexpected instantiate config: {config_node!r}')
def fake_dataloader(_dataset, *, shuffle, **_kwargs):
del shuffle, _kwargs
return _FakeLoader(
{
'observation.front': torch.zeros(1, 3, 2, 2),
'observation.state': torch.zeros(1, 4),
'action': torch.zeros(1, 2),
'action_is_pad': torch.zeros(1, 1, dtype=torch.bool),
},
length=1,
)
with tempfile.TemporaryDirectory() as tempdir:
previous_cwd = os.getcwd()
try:
os.chdir(tempdir)
with mock.patch.object(train_vla, 'instantiate', side_effect=fake_instantiate), \
mock.patch.object(train_vla, 'DataLoader', side_effect=fake_dataloader), \
mock.patch.object(train_vla, 'build_training_optimizer', return_value=_FakeOptimizer(cfg.train.lr)), \
mock.patch.object(train_vla, 'get_lr_schedule_with_warmup', return_value=_FakeScheduler()), \
mock.patch.object(train_vla, 'tqdm', side_effect=lambda iterable, **kwargs: _FakeProgressBar(iterable)), \
mock.patch.object(train_vla, '_init_swanlab', return_value=None), \
mock.patch.object(train_vla, '_finish_swanlab', return_value=None), \
mock.patch.object(train_vla.torch, 'save', return_value=None):
train_vla._run_training(cfg)
finally:
os.chdir(previous_cwd)
episode_indices_calls = [kwargs.get('episode_indices') for kwargs in instantiate_kwargs]
self.assertEqual(episode_indices_calls, [None, [101], [100]])
def test_action_mse_requires_explicit_val_episode_indices(self):
cfg = OmegaConf.create(
{
'agent': {
'normalization_type': 'min_max',
},
'data': {
'dataset_dir': 'unused',
'camera_names': ['front'],
},
'train': {
'batch_size': 2,
'lr': 1e-4,
'max_steps': 0,
'device': 'cpu',
'disable_cudnn': False,
'num_workers': 0,
'val_split': 0.25,
'val_episode_indices': None,
'action_mse_val_freq_epochs': 1,
'seed': 42,
'log_freq': 1,
'save_freq': 10,
'use_swanlab': False,
'rollout_val_freq_epochs': 0,
'rollout_validate_on_checkpoint': False,
'rollout_num_episodes': 1,
'warmup_steps': 1,
'scheduler_type': 'constant',
'min_lr': 1e-6,
'weight_decay': 1e-5,
'grad_clip': 1.0,
'pretrained_ckpt': None,
'resume_ckpt': None,
},
'eval': {
'ckpt_path': 'unused.pt',
'num_episodes': 1,
'headless': True,
'device': 'cpu',
'verbose_action': False,
},
'experiment': {},
}
)
def fake_instantiate(config_node, **_kwargs):
if config_node is cfg.data:
return _FakeDataset()
if config_node is cfg.agent:
return _FakeAgent()
raise AssertionError(f'unexpected instantiate config: {config_node!r}')
with tempfile.TemporaryDirectory() as tempdir:
previous_cwd = os.getcwd()
try:
os.chdir(tempdir)
with mock.patch.object(train_vla, 'instantiate', side_effect=fake_instantiate), \
mock.patch.object(train_vla, 'build_training_optimizer', return_value=_FakeOptimizer(cfg.train.lr)), \
mock.patch.object(train_vla, 'get_lr_schedule_with_warmup', return_value=_FakeScheduler()), \
mock.patch.object(train_vla, '_init_swanlab', return_value=None), \
mock.patch.object(train_vla, '_finish_swanlab', return_value=None):
with self.assertRaisesRegex(ValueError, 'val_episode_indices'):
train_vla._run_training(cfg)
finally:
os.chdir(previous_cwd)
def test_compute_action_mse_validation_masks_padding(self):
agent = _FakeAgent()
val_loader = _FakeLoader(
{
'observation.front': torch.zeros(1, 2, 3, 2, 2),
'observation.state': torch.zeros(1, 2, 4),
'action': torch.tensor([[[1.0, 1.0], [5.0, 5.0]]]),
'action_is_pad': torch.tensor([[False, True]]),
},
length=1,
)
mse = train_vla.compute_action_mse_validation(agent, val_loader, device='cpu')
self.assertEqual(mse, 0.0)
def test_eval_main_delegates_to_plain_run_eval_helper(self):
cfg = OmegaConf.create(
{
tests/test_train_vla_swanlab_logging.py
+70
View File
@@ -113,6 +113,20 @@ class FakeAgent(nn.Module):
def get_normalization_stats(self):
return {}
def predict_action_chunk(self, agent_input):
return agent_input['action']
class FakeEpisodeDataset(FakeDataset):
def __init__(self, episode_indices=None):
self.available_episode_indices = [100, 101]
self.episode_indices = None if episode_indices is None else list(episode_indices)
def __len__(self):
if self.episode_indices is None:
return 8
return 4 * len(self.episode_indices)
class FakeSwanLab:
def __init__(self, init_error=None, log_errors=None, finish_error=None, image_errors=None):
@@ -339,6 +353,8 @@ class TrainVLASwanLabLoggingTest(unittest.TestCase):
batch_size=2,
num_workers=0,
val_split=0.25,
val_episode_indices=None,
action_mse_val_freq_epochs=0,
seed=0,
lr=1e-3,
max_steps=2,
@@ -442,6 +458,8 @@ class TrainVLASwanLabLoggingTest(unittest.TestCase):
'batch_size': 2,
'num_workers': 0,
'val_split': 0.25,
'val_episode_indices': None,
'action_mse_val_freq_epochs': 0,
'seed': 0,
'lr': 1e-3,
'max_steps': 2,
@@ -539,6 +557,58 @@ class TrainVLASwanLabLoggingTest(unittest.TestCase):
self.assertEqual(fake_swanlab.finish_calls, 1)
def test_run_training_logs_action_mse_only_for_explicit_val_episode_validation(self):
module = self._load_train_vla_module()
run_training = self._get_run_training(module)
cfg = self._make_cfg()
cfg.train.val_split = 0.0
cfg.train.val_episode_indices = [100]
cfg.train.action_mse_val_freq_epochs = 1
cfg.train.max_steps = 4
cfg.train.save_freq = 100
cfg.train.rollout_val_freq_epochs = 0
agent = FakeAgent()
fake_swanlab = FakeSwanLab()
real_import_module = importlib.import_module
def fake_instantiate(config_node, **kwargs):
if config_node is cfg.data:
return FakeEpisodeDataset(kwargs.get('episode_indices'))
if config_node is cfg.agent:
return agent
raise AssertionError(f'unexpected instantiate config: {config_node!r}')
def fake_import_module(name, package=None):
if name == 'swanlab':
return fake_swanlab
return real_import_module(name, package)
def loader_factory(_dataset, *, shuffle, **_kwargs):
batch = {
'observation.front': torch.zeros(1, 2, 3, 2, 2),
'observation.state': torch.zeros(1, 2, 4),
'action': torch.zeros(1, 2, 2),
'action_is_pad': torch.zeros(1, 2, dtype=torch.bool),
}
return FakeLoader(batch)
with tempfile.TemporaryDirectory() as tempdir:
previous_cwd = os.getcwd()
try:
os.chdir(tempdir)
with mock.patch.object(module, 'instantiate', side_effect=fake_instantiate), \
mock.patch.object(module, 'DataLoader', side_effect=loader_factory), \
mock.patch.object(module, 'get_lr_schedule_with_warmup', return_value=FakeScheduler()), \
mock.patch.object(module, 'tqdm', side_effect=lambda iterable, **kwargs: FakeProgressBar(iterable)), \
mock.patch.object(module.torch, 'save', return_value=None), \
mock.patch.object(module.importlib, 'import_module', side_effect=fake_import_module):
run_training(cfg)
finally:
os.chdir(previous_cwd)
logged_keys = set().union(*(payload.keys() for payload, _step in fake_swanlab.log_calls))
self.assertIn('val/action_mse', logged_keys)
def test_run_training_warns_and_continues_when_swanlab_log_and_finish_fail(self):
module = self._load_train_vla_module()
run_training = self._get_run_training(module)