Brax

* Add boolean `wrap_env` to all brax `train` functions, which optionally wraps the env for training, or uses the env as is.
* Fix bug in PPO train to return loaded checkpoint when `num_timesteps` is 0.
* Add `layer_norm` to `make_q_network` and set `layer_norm` to `True` in `make_sace_networks` Q Network.
* Change PPO train function to return both value and policy network params, rather than just policy params.
* Merge https://github.com/google/brax/pull/561, adds grad norm clipping to PPO.
* Merge https://github.com/google/brax/issues/477, changes pusher vel damping.
* Merge https://github.com/google/brax/pull/558, adds `mocap_pos` and `mocap_quat` to render function.
* Merge https://github.com/google/brax/pull/559, allows for dictionary observations environment `State`.
* Merge https://github.com/google/brax/pull/562, which supports asymmetric actor-critic for PPO.
* Merge https://github.com/google/brax/pull/560, allows PPO from vision.

* Remove contact debugging from the viewer. This is a breaking change compared to older versions of brax. Old saved files will not render using the new viewer.
* Added `ctrl` as input to the `MJX pipeline.init`.
* Fixes to 513, 512, and 504.

* Modify `policy_params_fn` in the PPO implementation to take in the full model params. This can be used for checkpointing models.
* Add `restore_checkpoint_path` in PPO implementation.

- Add support for compressing json embedded in HTML output for large models.
- Remove legacy `dt` field in the brax System. We rely on MuJoCo's `opt.timestep` field instead.
- Add `act` to `pipeline.init` functions. See https://mujoco.readthedocs.io/en/stable/computation/index.html#physics-state.
- Updated basic APG algorithm 476, h/t Andrew-Luo1.

- Fix a bug in rendering capsules and cylinders with the wrong size.