Fix pickle env testing

Farama-Foundation · Apr 18, 2024 · ca84fb9 · ca84fb9
1 parent 6a8a267
commit ca84fb9
Showing 1 changed file with 9 additions and 3 deletions.
diff --git a/tests/envs/test_envs.py b/tests/envs/test_envs.py
@@ -149,11 +149,17 @@ def test_pickle_env(env: gym.Env):
     if env.metadata.get("jax", False):
         env = gym.wrappers.JaxToNumpy(env)
 
+    action = env.action_space.sample()
+
+    env_reset = env.reset(seed=123)
+    env_step = env.step(action)
+
     pickled_env = pickle.loads(pickle.dumps(env))
+    pickle_reset = pickled_env.reset(seed=123)
+    pickle_step = pickled_env.step(action)
 
-    data_equivalence(env.reset(), pickled_env.reset())
+    assert data_equivalence(env_reset, pickle_reset)
+    assert data_equivalence(env_step, pickle_step)
 
-    action = env.action_space.sample()
-    data_equivalence(env.step(action), pickled_env.step(action))
     env.close()
     pickled_env.close()