[Feat] test POMO's MS with more envs #102

ai4co · Nov 29, 2023 · d4acda9 · d4acda9
1 parent 91a76cd
commit d4acda9
Showing 1 changed file with 6 additions and 3 deletions.
diff --git a/tests/test_policy.py b/tests/test_policy.py
@@ -17,14 +17,17 @@ def test_base_policy(env_name, size=20, batch_size=2):
     assert out["reward"].shape == (batch_size,)
 
 
-@pytest.mark.parametrize("env_name", ["tsp", "cvrp", "pctsp", "spctsp"])
+@pytest.mark.parametrize(
+    "env_name", ["tsp", "cvrp", "pctsp", "spctsp", "sdvrp", "op", "pdp"]
+)
 def test_base_policy_multistart(env_name, size=20, batch_size=2):
     env, x = generate_env_data(env_name, size, batch_size)
     td = env.reset(x)
     policy = AutoregressivePolicy(env.name)
-    out = policy(td, env, decode_type="greedy_multistart", num_starts=size)
+    num_starts = size // 2 if env.name in ["pdp"] else size
+    out = policy(td, env, decode_type="greedy_multistart", num_starts=num_starts)
     assert out["reward"].shape == (
-        batch_size * size,
+        batch_size * num_starts,
     )  # to evaluate, we could just unbatchify