Add CurrentStep reward test

P6-Pool · Apr 11, 2024 · 55e943d · 55e943d
1 parent f8cb888
commit 55e943d
Showing 1 changed file with 10 additions and 0 deletions.
diff --git a/src/tests/utils/test_reward_functions.py b/src/tests/utils/test_reward_functions.py
@@ -264,6 +264,16 @@ def test_weights(self):
             1 / 5,
         )
 
+        reward = ConstantReward(weight=ConstantWeightCurrentStep, max_episode_steps=10)
+        reward.reset(table_state)
+        for _ in range(1, 3):
+            reward.get_reward(table_state, table_state, self.empty_action)
+
+        self.assertEqual(
+            reward.get_reward(table_state, table_state, self.empty_action),
+            1 / 3,
+        )
+
 
 if __name__ == "__main__":
     unittest.main()