Update for new pmap.

PiperOrigin-RevId: 815939422
Change-Id: Ife3b7328864884512183301d9685ad6d9bc69dcb

Author: btaba

brax/envs/__init__.py

@@ -32,12 +32,6 @@
 from brax.envs import walker2d
 from brax.envs.base import Env, PipelineEnv, State, Wrapper
 from brax.envs.wrappers import training
-import jax
-
-# NOTE(dsuo): Opt out of using the new `jax.pmap` implementation. The new
-# version is implemented using `jax.jit` and `jax.shard_map` and is causing
-# environment reset counts to be wrong.
-jax.config.update('jax_pmap_shmap_merge', False)
 
 _envs = {
     'ant': ant.Ant,