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

feat(sim): switch air insert task to socket peg

Browse Source
This commit is contained in:
Logic
2026-05-02 17:34:43 +08:00
parent 4c3646a3d5
commit 5c5cb299e9
16 changed files with 594 additions and 630 deletions
Download Patch File Download Diff File Expand all files Collapse all files

34
tests/test_eval_vla_headless.py
View File

@@ -114,7 +114,7 @@ class EvalVLAHeadlessTest(unittest.TestCase):
is_render=False,
control_freq=30,
is_interpolate=True,
cam_view="angle",
cam_view="left_side",
)
def test_camera_viewer_headless_updates_images_without_gui_calls(self):
@@ -123,11 +123,11 @@ class EvalVLAHeadlessTest(unittest.TestCase):
env.mj_data = object()
env.exit_flag = False
env.is_render = False
env.cam = "angle"
env.cam = "left_side"
env.r_vis = None
env.l_vis = None
env.top = None
env.angle = None
env.left_side = None
env.front = None
with mock.patch(
@@ -144,7 +144,7 @@ class EvalVLAHeadlessTest(unittest.TestCase):
self.assertIsNotNone(env.r_vis)
self.assertIsNotNone(env.l_vis)
self.assertIsNotNone(env.top)
self.assertIsNotNone(env.angle)
self.assertIsNotNone(env.left_side)
self.assertIsNotNone(env.front)
def test_eval_main_headless_skips_render_and_still_executes_policy(self):
@@ -254,19 +254,19 @@ class EvalVLAHeadlessTest(unittest.TestCase):
self.assertAlmostEqual(summary["avg_reward"], 3.75)
self.assertEqual(summary["num_episodes"], 2)
def test_run_eval_uses_air_insert_sampler_for_ring_bar_task(self):
def test_run_eval_uses_air_insert_sampler_for_socket_peg_task(self):
self.assertTrue(
hasattr(eval_vla, "sample_air_insert_ring_bar_state"),
"Expected eval_vla to expose the new ring/bar reset sampler",
hasattr(eval_vla, "sample_air_insert_socket_peg_state"),
"Expected eval_vla to expose the new socket/peg reset sampler",
)
fake_env = _FakeEnv()
fake_agent = _FakeAgent()
sampled_task_state = {
"ring_pos": np.array([-0.10, 0.80, 0.47], dtype=np.float32),
"ring_quat": np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32),
"bar_pos": np.array([0.10, 0.82, 0.47], dtype=np.float32),
"bar_quat": np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32),
"socket_pos": np.array([-0.10, 0.80, 0.47], dtype=np.float32),
"socket_quat": np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32),
"peg_pos": np.array([0.10, 0.82, 0.47], dtype=np.float32),
"peg_quat": np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32),
}
cfg = OmegaConf.create(
{
@@ -276,7 +276,7 @@ class EvalVLAHeadlessTest(unittest.TestCase):
"num_episodes": 1,
"max_timesteps": 1,
"device": "cpu",
"task_name": "sim_air_insert_ring_bar",
"task_name": "sim_air_insert_socket_peg",
"camera_names": ["front"],
"use_smoothing": False,
"smooth_alpha": 0.3,
@@ -296,12 +296,12 @@ class EvalVLAHeadlessTest(unittest.TestCase):
return_value=fake_env,
) as make_env, mock.patch.object(
eval_vla,
"sample_air_insert_ring_bar_state",
"sample_air_insert_socket_peg_state",
return_value=sampled_task_state,
) as ring_bar_sampler, mock.patch.object(
) as socket_peg_sampler, mock.patch.object(
eval_vla,
"sample_transfer_pose",
side_effect=AssertionError("sample_transfer_pose should not be used for sim_air_insert_ring_bar"),
side_effect=AssertionError("sample_transfer_pose should not be used for sim_air_insert_socket_peg"),
), mock.patch.object(
eval_vla,
"execute_policy_action",
@@ -312,8 +312,8 @@ class EvalVLAHeadlessTest(unittest.TestCase):
):
eval_vla._run_eval(cfg)
make_env.assert_called_once_with("sim_air_insert_ring_bar", headless=True)
ring_bar_sampler.assert_called_once_with()
make_env.assert_called_once_with("sim_air_insert_socket_peg", headless=True)
socket_peg_sampler.assert_called_once_with()
execute_policy_action.assert_called_once()
self.assertEqual(fake_env.reset_calls, [sampled_task_state])