From 9cd206a6875d08333de6333bfa480339a4646208 Mon Sep 17 00:00:00 2001 From: bridgesign Date: Wed, 27 Nov 2024 16:04:15 -0800 Subject: [PATCH] Reduce size of test --- brax/envs/wrappers/gym_test.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/brax/envs/wrappers/gym_test.py b/brax/envs/wrappers/gym_test.py index 2a1e01ff..29d286a0 100644 --- 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()