From fa40cc8c87764c7b8ae1052b75efbcc1371b3347 Mon Sep 17 00:00:00 2001 From: Jim Jiang Date: Tue, 16 Jul 2024 14:07:14 +0000 Subject: [PATCH] fix test_reset() function bug --- tests/test_methods.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/tests/test_methods.py b/tests/test_methods.py index 4a2d9f4..ed6eb02 100644 --- a/tests/test_methods.py +++ b/tests/test_methods.py @@ -65,13 +65,11 @@ def test_reset(): steps = env.Tmax # run the simulation util the for i in range(steps): - env.step(np.array([0,0, 0])) # set constant amount of trap every year - - curr_ob_space = env.observation_space + observation, reward, terminated, truncated, info = env.step(np.array([0,0, 0])) # set constant amount of trap every year # if final observation space is [-1, -1], set it to other obseravtion space for test - if curr_ob_space == np.array([-1,-1]): - curr_ob_space = np.array([1,1]) + if all(observation == np.array([-1,-1])): + observation = np.array([1,1]) # reset the obseravtion environment new_ob_space, new_info = env.reset()