Internal change

PiperOrigin-RevId: 703237189
Change-Id: I8f6d33e6b374333f0a9900fada70b6fc3c956b15

README.md

@@ -100,7 +100,7 @@ If you would like to reference Brax in a publication, please use:
   author = {C. Daniel Freeman and Erik Frey and Anton Raichuk and Sertan Girgin and Igor Mordatch and Olivier Bachem},
   title = {Brax - A Differentiable Physics Engine for Large Scale Rigid Body Simulation},
   url = {http://github.com/google/brax},
-  version = {0.11.0},
+  version = {0.12.0},
   year = {2021},
 }
 ```

brax/__init__.py

@@ -14,7 +14,7 @@
 
 """Import top-level classes and functions here for encapsulation/clarity."""
 
-__version__ = '0.11.0'
+__version__ = '0.12.0'
 
 from brax.base import Motion
 from brax.base import State

brax/io/json.py

@@ -16,12 +16,11 @@
 """Saves a system config and trajectory as json."""
 
 import json
-from typing import List, Optional, Text, Tuple
+from typing import List, Text, Tuple
 
 from brax.base import State, System
 from etils import epath
 import jax
-import jax.numpy as jp
 import mujoco
 import numpy as np
 
@@ -128,7 +127,7 @@ def dumps(sys: System, states: List[State]) -> Text:
   for id_ in range(sys.ngeom):
     link_idx = sys.geom_bodyid[id_] - 1
 
-    rgba = sys.mj_model.geom_rgba[id_]
+    rgba = np.array(sys.mj_model.geom_rgba[id_])
     if (rgba == [0.5, 0.5, 0.5, 1.0]).all():
       # convert the default mjcf color to brax default color
       rgba = np.array([0.4, 0.33, 0.26, 1.0])

docs/release-notes/next-release.md

@@ -1,12 +1 @@
 # Brax Release Notes
-
-* 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.

docs/release-notes/v0.12.0.md

@@ -0,0 +1,12 @@
+# Brax v0.12.0 Release Notes
+
+* 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.

setup.py

@@ -24,7 +24,7 @@
 
 setup(
     name="brax",
-    version="0.11.0",
+    version="0.12.0",
     description="A differentiable physics engine written in JAX.",
     author="Brax Authors",
     author_email="no-reply@google.com",