Keep all ExperimentBuilder tests in one place

thu-ml · Oct 24, 2023 · 58466eb · 58466eb
1 parent b5a8915
commit 58466eb
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 42 deletions.
diff --git a/test/highlevel/test_discrete.py b/test/highlevel/test_discrete.py
diff --git a/test/highlevel/test_continuous.py → test/highlevel/test_experiment_builder.py b/test/highlevel/test_continuous.py → test/highlevel/test_experiment_builder.py
@@ -1,12 +1,15 @@
-from test.highlevel.env_factory import ContinuousTestEnvFactory
+from test.highlevel.env_factory import ContinuousTestEnvFactory, DiscreteTestEnvFactory
 
 import pytest
 
 from tianshou.highlevel.config import SamplingConfig
 from tianshou.highlevel.experiment import (
     A2CExperimentBuilder,
     DDPGExperimentBuilder,
+    DiscreteSACExperimentBuilder,
+    DQNExperimentBuilder,
     ExperimentConfig,
+    IQNExperimentBuilder,
     PGExperimentBuilder,
     PPOExperimentBuilder,
     REDQExperimentBuilder,
@@ -47,3 +50,31 @@ def test_experiment_builder_continuous_default_params(builder_cls):
     experiment = builder.build()
     experiment.run("test")
     print(experiment)
+
+
+@pytest.mark.parametrize(
+    "builder_cls",
+    [
+        PPOExperimentBuilder,
+        A2CExperimentBuilder,
+        DQNExperimentBuilder,
+        DiscreteSACExperimentBuilder,
+        IQNExperimentBuilder,
+    ],
+)
+def test_experiment_builder_discrete_default_params(builder_cls):
+    env_factory = DiscreteTestEnvFactory()
+    sampling_config = SamplingConfig(
+        num_epochs=1,
+        step_per_epoch=100,
+        num_train_envs=2,
+        num_test_envs=2,
+    )
+    builder = builder_cls(
+        experiment_config=ExperimentConfig(persistence_enabled=False),
+        env_factory=env_factory,
+        sampling_config=sampling_config,
+    )
+    experiment = builder.build()
+    experiment.run("test")
+    print(experiment)