feat(sim): save air-insert and rollout validation updates
This commit is contained in:
Logic
@@ -102,10 +102,8 @@ 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)
|
||||
|
||||
@@ -135,8 +133,6 @@ 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,
|
||||
@@ -180,14 +176,12 @@ 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]))
|
||||
@@ -224,120 +218,267 @@ 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,
|
||||
def test_run_eval_parallel_rejects_trajectory_and_video_exports(self):
|
||||
unsupported_flags = [
|
||||
"record_video",
|
||||
"save_trajectory",
|
||||
"save_trajectory_npz",
|
||||
]
|
||||
fake_agent = _FakeAgent(actions)
|
||||
fake_env = _FakeEnv()
|
||||
|
||||
for flag_name in unsupported_flags:
|
||||
with self.subTest(flag_name=flag_name):
|
||||
cfg = OmegaConf.create(
|
||||
{
|
||||
"agent": {},
|
||||
"eval": {
|
||||
"ckpt_path": "checkpoints/vla_model_best.pt",
|
||||
"num_episodes": 2,
|
||||
"num_workers": 2,
|
||||
"max_timesteps": 1,
|
||||
"device": "cpu",
|
||||
"task_name": "sim_transfer",
|
||||
"camera_names": ["front"],
|
||||
"use_smoothing": False,
|
||||
"smooth_alpha": 0.3,
|
||||
"verbose_action": False,
|
||||
"headless": True,
|
||||
"save_artifacts": True,
|
||||
flag_name: True,
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
with self.assertRaisesRegex(ValueError, flag_name):
|
||||
eval_vla._run_eval_parallel(cfg)
|
||||
|
||||
def test_run_eval_parallel_writes_merged_summary_timing_and_worker_dirs(self):
|
||||
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,
|
||||
"agent": {},
|
||||
"eval": {
|
||||
"ckpt_path": "checkpoints/vla_model_best.pt",
|
||||
"num_episodes": 3,
|
||||
"num_workers": 2,
|
||||
"max_timesteps": 1,
|
||||
"device": "cpu",
|
||||
"task_name": "sim_transfer",
|
||||
"camera_names": ["front"],
|
||||
"use_smoothing": False,
|
||||
"smooth_alpha": 0.3,
|
||||
"verbose_action": False,
|
||||
"headless": True,
|
||||
"artifact_dir": tmpdir,
|
||||
"save_summary_json": True,
|
||||
"save_timing": True,
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
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(
|
||||
def fake_run_spawn_jobs(payloads, max_workers, worker_fn):
|
||||
del max_workers, worker_fn
|
||||
return [
|
||||
{
|
||||
'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)
|
||||
"episodes": [
|
||||
{
|
||||
"episode_index": 2,
|
||||
"episode_reward": 3.0,
|
||||
"episode_max_reward": 3.0,
|
||||
"inference_fps": 30.0,
|
||||
"control_fps": 15.0,
|
||||
}
|
||||
],
|
||||
"_merge_state": {
|
||||
"obs_read_time_ms": [3.0],
|
||||
"preprocess_time_ms": [1.0],
|
||||
"inference_time_ms": [2.0],
|
||||
"env_step_time_ms": [4.0],
|
||||
"total_time_ms": [5.0],
|
||||
"model_forward_flags": [True],
|
||||
},
|
||||
},
|
||||
{
|
||||
"episodes": [
|
||||
{
|
||||
"episode_index": 1,
|
||||
"episode_reward": 2.0,
|
||||
"episode_max_reward": 2.0,
|
||||
"inference_fps": 20.0,
|
||||
"control_fps": 10.0,
|
||||
},
|
||||
{
|
||||
"episode_index": 0,
|
||||
"episode_reward": 1.0,
|
||||
"episode_max_reward": 1.0,
|
||||
"inference_fps": 10.0,
|
||||
"control_fps": 5.0,
|
||||
},
|
||||
],
|
||||
"_merge_state": {
|
||||
"obs_read_time_ms": [1.0, 2.0],
|
||||
"preprocess_time_ms": [1.0, 1.0],
|
||||
"inference_time_ms": [2.0, 2.0],
|
||||
"env_step_time_ms": [4.0, 4.0],
|
||||
"total_time_ms": [5.0, 5.0],
|
||||
"model_forward_flags": [False, True],
|
||||
},
|
||||
},
|
||||
]
|
||||
|
||||
with mock.patch.object(
|
||||
eval_vla,
|
||||
'load_checkpoint',
|
||||
return_value=(fake_agent, None),
|
||||
"sample_transfer_pose",
|
||||
side_effect=[
|
||||
np.array([0.1, 0.2, 0.3], dtype=np.float32),
|
||||
np.array([0.4, 0.5, 0.6], dtype=np.float32),
|
||||
np.array([0.7, 0.8, 0.9], dtype=np.float32),
|
||||
],
|
||||
), 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)
|
||||
"_run_spawn_jobs",
|
||||
side_effect=fake_run_spawn_jobs,
|
||||
):
|
||||
summary = eval_vla._run_eval_parallel(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'],
|
||||
summary_path = Path(tmpdir) / "rollout_summary.json"
|
||||
timing_path = Path(tmpdir) / "timing.json"
|
||||
worker_00_dir = Path(tmpdir) / "workers" / "worker_00"
|
||||
worker_01_dir = Path(tmpdir) / "workers" / "worker_01"
|
||||
|
||||
self.assertTrue(summary_path.exists())
|
||||
self.assertTrue(timing_path.exists())
|
||||
self.assertTrue(worker_00_dir.is_dir())
|
||||
self.assertTrue(worker_01_dir.is_dir())
|
||||
self.assertEqual(summary["episode_rewards"], [1.0, 2.0, 3.0])
|
||||
|
||||
with summary_path.open("r", encoding="utf-8") as fh:
|
||||
saved_summary = json.load(fh)
|
||||
with timing_path.open("r", encoding="utf-8") as fh:
|
||||
saved_timing = json.load(fh)
|
||||
|
||||
self.assertEqual(saved_summary["episode_rewards"], [1.0, 2.0, 3.0])
|
||||
self.assertEqual(saved_summary["artifact_dir"], tmpdir)
|
||||
self.assertEqual(saved_timing["count"], 3)
|
||||
self.assertEqual(saved_timing["model_forward_count"], 2)
|
||||
|
||||
def test_run_eval_parallel_cuda_writes_merged_summary_timing_and_worker_dirs(self):
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
cfg = OmegaConf.create(
|
||||
{
|
||||
"agent": {},
|
||||
"eval": {
|
||||
"ckpt_path": "checkpoints/vla_model_best.pt",
|
||||
"num_episodes": 3,
|
||||
"num_workers": 2,
|
||||
"cuda_devices": [0],
|
||||
"max_timesteps": 1,
|
||||
"device": "cuda",
|
||||
"task_name": "sim_transfer",
|
||||
"camera_names": ["front"],
|
||||
"use_smoothing": False,
|
||||
"smooth_alpha": 0.3,
|
||||
"verbose_action": False,
|
||||
"headless": True,
|
||||
"artifact_dir": tmpdir,
|
||||
"save_summary_json": True,
|
||||
"save_timing": True,
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
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))
|
||||
def fake_run_cuda_parallel_processes(server_payloads, worker_payloads):
|
||||
self.assertEqual(len(server_payloads), 1)
|
||||
self.assertEqual(server_payloads[0]["device_index"], 0)
|
||||
self.assertEqual([payload["server_index"] for payload in worker_payloads], [0, 0])
|
||||
return [
|
||||
{
|
||||
"episodes": [
|
||||
{
|
||||
"episode_index": 2,
|
||||
"episode_reward": 3.0,
|
||||
"episode_max_reward": 3.0,
|
||||
"inference_fps": 30.0,
|
||||
"control_fps": 15.0,
|
||||
}
|
||||
],
|
||||
"_merge_state": {
|
||||
"obs_read_time_ms": [3.0],
|
||||
"preprocess_time_ms": [1.0],
|
||||
"inference_time_ms": [2.0],
|
||||
"env_step_time_ms": [4.0],
|
||||
"total_time_ms": [5.0],
|
||||
"model_forward_flags": [True],
|
||||
},
|
||||
},
|
||||
{
|
||||
"episodes": [
|
||||
{
|
||||
"episode_index": 1,
|
||||
"episode_reward": 2.0,
|
||||
"episode_max_reward": 2.0,
|
||||
"inference_fps": 20.0,
|
||||
"control_fps": 10.0,
|
||||
},
|
||||
{
|
||||
"episode_index": 0,
|
||||
"episode_reward": 1.0,
|
||||
"episode_max_reward": 1.0,
|
||||
"inference_fps": 10.0,
|
||||
"control_fps": 5.0,
|
||||
},
|
||||
],
|
||||
"_merge_state": {
|
||||
"obs_read_time_ms": [1.0, 2.0],
|
||||
"preprocess_time_ms": [1.0, 1.0],
|
||||
"inference_time_ms": [2.0, 2.0],
|
||||
"env_step_time_ms": [4.0, 4.0],
|
||||
"total_time_ms": [5.0, 5.0],
|
||||
"model_forward_flags": [False, True],
|
||||
},
|
||||
},
|
||||
]
|
||||
|
||||
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])
|
||||
with mock.patch.object(
|
||||
eval_vla,
|
||||
"sample_transfer_pose",
|
||||
side_effect=[
|
||||
np.array([0.1, 0.2, 0.3], dtype=np.float32),
|
||||
np.array([0.4, 0.5, 0.6], dtype=np.float32),
|
||||
np.array([0.7, 0.8, 0.9], dtype=np.float32),
|
||||
],
|
||||
), mock.patch.object(
|
||||
eval_vla,
|
||||
"_run_cuda_parallel_processes",
|
||||
side_effect=fake_run_cuda_parallel_processes,
|
||||
create=True,
|
||||
):
|
||||
summary = eval_vla._run_eval_parallel_cuda(cfg)
|
||||
|
||||
summary_path = Path(tmpdir) / "rollout_summary.json"
|
||||
timing_path = Path(tmpdir) / "timing.json"
|
||||
worker_00_dir = Path(tmpdir) / "workers" / "worker_00"
|
||||
worker_01_dir = Path(tmpdir) / "workers" / "worker_01"
|
||||
|
||||
self.assertTrue(summary_path.exists())
|
||||
self.assertTrue(timing_path.exists())
|
||||
self.assertTrue(worker_00_dir.is_dir())
|
||||
self.assertTrue(worker_01_dir.is_dir())
|
||||
self.assertEqual(summary["episode_rewards"], [1.0, 2.0, 3.0])
|
||||
|
||||
with summary_path.open("r", encoding="utf-8") as fh:
|
||||
saved_summary = json.load(fh)
|
||||
with timing_path.open("r", encoding="utf-8") as fh:
|
||||
saved_timing = json.load(fh)
|
||||
|
||||
self.assertEqual(saved_summary["episode_rewards"], [1.0, 2.0, 3.0])
|
||||
self.assertEqual(saved_summary["artifact_dir"], tmpdir)
|
||||
self.assertEqual(saved_timing["count"], 3)
|
||||
self.assertEqual(saved_timing["model_forward_count"], 2)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
|
||||
Reference in New Issue
Block a user