[brax] Explicitly set jax_pmap_shmap_merge=False.

Part of a larger clean-up to prepare for migrating `jax.pmap` users to `jax.shard_map` or the new `jax.pmap` that's implemented using `jax.shard_map`.

PiperOrigin-RevId: 811785897
Change-Id: I32b92af43d1bcf2605e15582c80c8e4f8df2d75d

Author: danielsuo

brax/envs/__init__.py

@@ -32,6 +32,12 @@
 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,