Added a VectorizedGymWrapper

brax/envs/__init__.py

@@ -15,7 +15,7 @@
 """Some example environments to help get started quickly with brax."""
 
 import functools
-from typing import Callable
+from typing import Callable, Union, overload
 
 import gym
 import brax
@@ -54,6 +54,21 @@ def create_fn(env_name: str, **kwargs) -> Callable[..., Env]:
   return functools.partial(create, env_name, **kwargs)
 
 
-def create_gym_env(env_name: str, **kwargs) -> gym.Env:
-  """Creates a Gym Env with a specified brax system."""
-  return wrappers.GymWrapper(create(env_name, **kwargs))
+@overload
+def create_gym_env(env_name: str, seed: int = 0, **kwargs) -> wrappers.GymWrapper:
+    ...
+
+
+@overload
+def create_gym_env(env_name: str, batch_size: int, seed: int = 0, **kwargs) -> wrappers.VectorGymWrapper:
+    ...
+
+
+def create_gym_env(
+    env_name: str, batch_size: int = 0, seed: int = 0, **kwargs
+) -> Union[wrappers.GymWrapper, wrappers.VectorGymWrapper]:
+    if batch_size:
+        return wrappers.VectorGymWrapper(create(env_name=env_name, batch_size=batch_size, **kwargs), seed=seed)
+    else:
+        return wrappers.GymWrapper(create(env_name=env_name, **kwargs), seed=seed)
+

brax/envs/wrappers.py

@@ -16,6 +16,7 @@
 
 import gym
 from gym import spaces
+from gym.vector.utils import batch_space
 import jax
 import numpy as np
 from brax.envs import env
@@ -26,7 +27,7 @@ class GymWrapper(gym.Env):
 
   def __init__(self, environment: env.Env, seed: int = 0):
     self._environment = environment
-    self._key = jax.random.PRNGKey(seed)
+    self.seed(seed)
 
     # action_space = None
     obs_high = np.inf * np.ones(self._environment.observation_size)
@@ -55,3 +56,49 @@ def reset(self):
   def step(self, action):
     self._state, obs, reward, done = self._step(self._state, action)
     return obs, reward, done, {}
+
+  def seed(self, seed: int = 0):
+    self._key = jax.random.PRNGKey(seed)
+
+
+class VectorGymWrapper(gym.vector.VectorEnv):
+  """A wrapper that converts batched Brax Env to one that follows Gym VectorEnv API."""
+
+  def __init__(self, environment: env.Env, seed: int = 0):
+    self._environment = environment
+    assert self._environment.batch_size  # Make sure underlying environment is batched
+
+    self.num_envs = self._environment.batch_size
+    self._key_size = self.num_envs + 1
+    self.seed(seed)
+
+    obs_high = np.inf * np.ones(self._environment.observation_size)
+    self.single_observation_space = spaces.Box(-obs_high, obs_high, dtype=np.float32)
+    self.observation_space = batch_space(self.single_observation_space, self.num_envs)
+
+    action_high = np.ones(self._environment.action_size)
+    self.single_action_space = spaces.Box(-action_high, action_high, dtype=np.float32)
+    self.action_space = batch_space(self.single_action_space, self.num_envs)
+    self._state = None
+
+    def reset(key):
+      keys = jax.random.split(key, self._key_size)
+      state = self._environment.reset(keys[1:])
+      return state, state.obs, keys[0]
+    self._reset = jax.jit(reset)
+
+    def step(state, action):
+      state = self._environment.step(state, action)
+      return state, state.obs, state.reward, state.done
+    self._step = jax.jit(step, backend='cpu')
+
+  def reset(self):
+    self._state, obs, self._key = self._reset(self._key)
+    return obs
+
+  def step(self, action):
+    self._state, obs, reward, done = self._step(self._state, action)
+    return obs, reward, done, {}
+
+  def seed(self, seed: int = 0):
+    self._key = jax.random.PRNGKey(seed)