[JAX] Replace jnp.clip(..., a_min=..., a_max=...) with jnp.clip(..., min=..., max=...).

[pmap] In-line definitions of `jax.device_put_sharded` and `jax.device_put_replicated`.

Prepare mujoco playground for release.

Update Brax version to 0.14.2.

Use GAE returns as the target for the value function loss in PPO.

Add distributional critic to PPO and enhance observation normalization.

Update hypers for cartpole vision.

Fix Brax ES argsort failure by resharding weights before fitness shaping

Add CNN kernel init and spatial softmax (h/t zakka)

Update Python CI versions (#663)

Import Brax PR #662: Vision update

Import Brax PR #662: Vision update

Fix Brax external tests by updating jax.sharding.PartitionSpec

Vision update (#662)

[pmap] Remove `jax.config.pmap_shmap_merge`.