VectorEnv gym visulization correction (#549)

* VectorEnv gym visulization correction

This is a solution in context of #535

* Added tests. Allow width-height configuration

* Reduce batch size

* Reduce size of test

* Change render to list pipeline states before conversion

* Update gym_test.py

* Update gym_test.py

---------

Co-authored-by: btaba <btaba@google.com>

Author: bridgesign

brax/envs/wrappers/gym.py

@@ -77,12 +77,12 @@ def step(self, action):
   def seed(self, seed: int = 0):
     self._key = jax.random.PRNGKey(seed)
 
-  def render(self, mode='human'):
+  def render(self, mode='human', width=256, height=256):
     if mode == 'rgb_array':
       sys, state = self._env.sys, self._state
       if state is None:
         raise RuntimeError('must call reset or step before rendering')
-      return image.render_array(sys, state.pipeline_state, 256, 256)
+      return image.render_array(sys, state.pipeline_state, width=width, height=height)
     else:
       return super().render(mode=mode)  # just raise an exception
 
@@ -143,11 +143,12 @@ def step(self, action):
   def seed(self, seed: int = 0):
     self._key = jax.random.PRNGKey(seed)
 
-  def render(self, mode='human'):
+  def render(self, mode='human', width=256, height=256):
     if mode == 'rgb_array':
       sys, state = self._env.sys, self._state
       if state is None:
         raise RuntimeError('must call reset or step before rendering')
-      return image.render_array(sys, state.pipeline_state.take(0), 256, 256)
+      state_list = [state.take(i).pipeline_state for i in range(self.num_envs)]
+      return np.stack(image.render_array(sys, state_list, width=width, height=height))
     else:
       return super().render(mode=mode)  # just raise an exception

brax/envs/wrappers/gym_test.py

@@ -46,7 +46,7 @@ def test_vector_action_space(self):
     np.testing.assert_array_equal(
         env.action_space.high,
         np.tile(base_env.sys.actuator.ctrl_range[:, 1], [256, 1]),
-    )
+    )  
 
 
 if __name__ == '__main__':