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

feat: add rollout trajectory image artifacts and swanlab logging

Browse Source
This commit is contained in:
Logic
2026-04-03 09:39:16 +08:00
parent 48f0eb8dd0
commit 0586a6e6c7
8 changed files with 626 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

116
tests/test_eval_vla_rollout_artifacts.py
View File

@@ -102,8 +102,10 @@ class EvalVLARolloutArtifactsTest(unittest.TestCase):
self.assertIn('artifact_dir', eval_cfg)
self.assertFalse(eval_cfg.save_summary_json)
self.assertFalse(eval_cfg.save_trajectory_npz)
self.assertFalse(eval_cfg.save_trajectory_image)
self.assertFalse(eval_cfg.record_video)
self.assertIsNone(eval_cfg.artifact_dir)
self.assertIsNone(eval_cfg.trajectory_image_camera_name)
self.assertIsNone(eval_cfg.video_camera_name)
self.assertEqual(eval_cfg.video_fps, 30)
@@ -133,6 +135,8 @@ class EvalVLARolloutArtifactsTest(unittest.TestCase):
'artifact_dir': tmpdir,
'save_summary_json': True,
'save_trajectory_npz': True,
'save_trajectory_image': True,
'trajectory_image_camera_name': 'front',
'record_video': True,
'video_camera_name': 'front',
'video_fps': 12,
@@ -176,12 +180,14 @@ class EvalVLARolloutArtifactsTest(unittest.TestCase):
trajectory_path = Path(artifacts['trajectory_npz'])
summary_path = Path(artifacts['summary_json'])
video_path = Path(artifacts['video_mp4'])
trajectory_image_path = Path(summary['episodes'][0]['artifact_paths']['trajectory_image'])
self.assertEqual(Path(artifacts['output_dir']), Path(tmpdir))
self.assertEqual(artifacts['video_camera_name'], 'front')
self.assertTrue(trajectory_path.exists())
self.assertTrue(summary_path.exists())
self.assertTrue(video_path.exists())
self.assertTrue(trajectory_image_path.exists())
rollout_npz = np.load(trajectory_path)
np.testing.assert_array_equal(rollout_npz['episode_index'], np.array([0, 0]))
@@ -218,11 +224,121 @@ class EvalVLARolloutArtifactsTest(unittest.TestCase):
saved_summary = json.load(fh)
self.assertEqual(saved_summary['artifacts']['trajectory_npz'], str(trajectory_path))
self.assertEqual(saved_summary['artifacts']['video_mp4'], str(video_path))
self.assertEqual(
saved_summary['episodes'][0]['artifact_paths']['trajectory_image'],
str(trajectory_image_path),
)
self.assertEqual(saved_summary['episode_rewards'], [3.0])
self.assertAlmostEqual(summary['avg_reward'], 3.0)
self.assertIn('avg_obs_read_time_ms', summary)
self.assertIn('avg_env_step_time_ms', summary)
def test_run_eval_exports_front_trajectory_images_without_video_dependency(self):
actions = [
np.arange(16, dtype=np.float32),
np.arange(16, dtype=np.float32) + 10.0,
np.arange(16, dtype=np.float32) + 100.0,
np.arange(16, dtype=np.float32) + 110.0,
]
fake_agent = _FakeAgent(actions)
fake_env = _FakeEnv()
with tempfile.TemporaryDirectory() as tmpdir:
cfg = OmegaConf.create(
{
'agent': {},
'eval': {
'ckpt_path': 'checkpoints/vla_model_best.pt',
'num_episodes': 2,
'max_timesteps': 2,
'device': 'cpu',
'task_name': 'sim_transfer',
'camera_names': ['top', 'front'],
'use_smoothing': True,
'smooth_alpha': 0.5,
'verbose_action': False,
'headless': True,
'artifact_dir': tmpdir,
'save_trajectory_image': True,
'record_video': False,
},
}
)
trajectory_image_calls = []
def fake_save_rollout_trajectory_image(
env,
output_path,
raw_actions,
camera_name,
*,
line_radius=0.004,
max_markers=1500,
):
del env, line_radius, max_markers
trajectory_image_calls.append(
{
'output_path': output_path,
'camera_name': camera_name,
'raw_actions': [np.array(action, copy=True) for action in raw_actions],
}
)
if output_path is None:
return None
output_path = Path(output_path)
output_path.parent.mkdir(parents=True, exist_ok=True)
output_path.write_bytes(b'fake-png')
return str(output_path)
with mock.patch.object(
eval_vla,
'load_checkpoint',
return_value=(fake_agent, None),
), mock.patch.object(
eval_vla,
'make_sim_env',
return_value=fake_env,
), mock.patch.object(
eval_vla,
'sample_transfer_pose',
return_value=np.array([0.1, 0.2, 0.3], dtype=np.float32),
), mock.patch.object(
eval_vla,
'tqdm',
side_effect=lambda iterable, **kwargs: iterable,
), mock.patch.object(
eval_vla,
'_save_rollout_trajectory_image',
side_effect=fake_save_rollout_trajectory_image,
) as save_trajectory_image_mock, mock.patch.object(
eval_vla,
'_open_video_writer',
) as open_video_writer_mock:
summary = eval_vla._run_eval(cfg)
self.assertEqual(save_trajectory_image_mock.call_count, 2)
open_video_writer_mock.assert_not_called()
self.assertIsNone(summary['artifacts']['video_mp4'])
self.assertEqual(summary['artifacts']['trajectory_image_camera_name'], 'front')
self.assertEqual(
[call['camera_name'] for call in trajectory_image_calls],
['front', 'front'],
)
first_episode_path = Path(summary['episodes'][0]['artifact_paths']['trajectory_image'])
second_episode_path = Path(summary['episodes'][1]['artifact_paths']['trajectory_image'])
self.assertTrue(first_episode_path.exists())
self.assertTrue(second_episode_path.exists())
self.assertNotEqual(first_episode_path, second_episode_path)
self.assertEqual(first_episode_path.parent, Path(tmpdir))
self.assertEqual(second_episode_path.parent, Path(tmpdir))
np.testing.assert_array_equal(trajectory_image_calls[0]['raw_actions'][0], actions[0])
np.testing.assert_array_equal(trajectory_image_calls[0]['raw_actions'][1], actions[1])
np.testing.assert_array_equal(trajectory_image_calls[1]['raw_actions'][0], actions[2])
np.testing.assert_array_equal(trajectory_image_calls[1]['raw_actions'][1], actions[3])
if __name__ == '__main__':
unittest.main()