Reduce size of test

google · Nov 28, 2024 · 9cd206a · 9cd206a
1 parent 2928965
commit 9cd206a
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/brax/envs/wrappers/gym_test.py b/brax/envs/wrappers/gym_test.py
@@ -55,10 +55,10 @@ def test_render(self):
   def test_vector_render(self):
     """Tests rendering in the VectorGymWrapper."""
     base_env = envs.create('pusher')
-    env = gym.VectorGymWrapper(training.VmapWrapper(base_env, batch_size=16))
+    env = gym.VectorGymWrapper(training.VmapWrapper(base_env, batch_size=2))
     env.reset()
-    img = env.render(mode='rgb_array', width=250, height=236)
-    self.assertEqual(img.shape, (16, 236, 250, 3))
+    img = env.render(mode='rgb_array', width=128, height=128)
+    self.assertEqual(img.shape, (2, 128, 128, 3))
 
 if __name__ == '__main__':
   absltest.main()