wip: format

opendilab · Mar 2, 2023 · 693d148 · 693d148
1 parent a8b710d
commit 693d148
Showing 1 changed file with 26 additions and 23 deletions.
diff --git a/ding/model/template/tests/test_pdqn.py b/ding/model/template/tests/test_pdqn.py
@@ -4,36 +4,39 @@
 
 from ding.model.template import PDQN
 
+action_args_shape_values = [1, 5]
+
 
 @pytest.mark.unittest
 class TestPQQN:
 
-    def test_dqn(self):
-        for action_args_shape in (1, 5):
-            T, B = 3, 4
-            obs_shape = (4, )
-            act_shape = EasyDict({'action_type_shape': (3, ), 'action_args_shape': (action_args_shape, )})
-            if isinstance(obs_shape, int):
-                cont_inputs = torch.randn(B, obs_shape)
-            else:
-                cont_inputs = torch.randn(B, *obs_shape)
-            model = PDQN(obs_shape, act_shape)
-            cont_outputs = model.forward(cont_inputs, mode='compute_continuous')
-            assert isinstance(cont_outputs, dict)
-            dis_inputs = {'state': cont_inputs, 'action_args': cont_outputs['action_args']}
-            dis_outputs = model.forward(dis_inputs, mode='compute_discrete')
-            assert isinstance(dis_outputs, dict)
-            if isinstance(act_shape['action_type_shape'], int):
-                assert dis_outputs['logit'].shape == (B, act_shape.action_type_shape)
-            elif len(act_shape['action_type_shape']) == 1:
-                assert dis_outputs['logit'].shape == (B, *act_shape.action_type_shape)
-            else:
-                for i, s in enumerate(act_shape):
-                    assert dis_outputs['logit'][i].shape == (B, s)
+    @pytest.mark.unittest
+    @pytest.mark.parametrize('action_type_shape', action_args_shape_values)
+    def test_dqn(self, action_type_shape):
+        T, B = 3, 4
+        obs_shape = (4,)
+        act_shape = EasyDict({'action_type_shape': (3,), 'action_args_shape': (action_type_shape,)})
+        if isinstance(obs_shape, int):
+            cont_inputs = torch.randn(B, obs_shape)
+        else:
+            cont_inputs = torch.randn(B, *obs_shape)
+        model = PDQN(obs_shape, act_shape)
+        cont_outputs = model.forward(cont_inputs, mode='compute_continuous')
+        assert isinstance(cont_outputs, dict)
+        dis_inputs = {'state': cont_inputs, 'action_args': cont_outputs['action_args']}
+        dis_outputs = model.forward(dis_inputs, mode='compute_discrete')
+        assert isinstance(dis_outputs, dict)
+        if isinstance(act_shape['action_type_shape'], int):
+            assert dis_outputs['logit'].shape == (B, act_shape.action_type_shape)
+        elif len(act_shape['action_type_shape']) == 1:
+            assert dis_outputs['logit'].shape == (B, *act_shape.action_type_shape)
+        else:
+            for i, s in enumerate(act_shape):
+                assert dis_outputs['logit'][i].shape == (B, s)
 
     def test_mdqn(self):
         T, B = 3, 4
-        obs_shape = (4, )
+        obs_shape = (4,)
         act_shape = EasyDict({'action_type_shape': 3, 'action_args_shape': 5})
         if isinstance(obs_shape, int):
             cont_inputs = torch.randn(B, obs_shape)