Internal change

PiperOrigin-RevId: 747504144
Change-Id: Ia6f000e46f84e5f62793dbf7ac347a5ca8d8f540

Author: Brax Team

MANIFEST.in

@@ -1,3 +1,7 @@
+prune docs
+prune notebooks
+prune datasets
+
 include brax/envs/assets/*.xml
 recursive-include brax/experimental/barkour *.csv *.stl *.xml
 recursive-include brax/test_data *.xml *.stl *.obj *.urdf

brax/training/agents/ppo/train.py

@@ -423,13 +423,6 @@ def train(
       progress_fn=progress_fn,
   )
 
-  ckpt_config = checkpoint.network_config(
-      observation_size=obs_shape,
-      action_size=env.action_size,
-      normalize_observations=normalize_observations,
-      network_factory=network_factory,
-  )
-
   def minibatch_step(
       carry,
       data: types.Transition,
@@ -713,6 +706,12 @@ def training_epoch_with_timing(
     policy_params_fn(current_step, make_policy, params)
 
     if save_checkpoint_path is not None:
+      ckpt_config = checkpoint.network_config(
+          observation_size=obs_shape,
+          action_size=env.action_size,
+          normalize_observations=normalize_observations,
+          network_factory=network_factory,
+      )
       checkpoint.save(
           save_checkpoint_path, current_step, params, ckpt_config
       )

brax/training/agents/sac/train.py

@@ -260,13 +260,6 @@ def train(
       actor_loss, policy_optimizer, pmap_axis_name=_PMAP_AXIS_NAME
   )
 
-  ckpt_config = checkpoint.network_config(
-      observation_size=obs_size,
-      action_size=env.action_size,
-      normalize_observations=normalize_observations,
-      network_factory=network_factory,
-  )
-
   def sgd_step(
       carry: Tuple[TrainingState, PRNGKey], transitions: Transition
   ) -> Tuple[Tuple[TrainingState, PRNGKey], Metrics]:
@@ -585,6 +578,12 @@ def training_epoch_with_timing(
         params = _unpmap(
             (training_state.normalizer_params, training_state.policy_params)
         )
+        ckpt_config = checkpoint.network_config(
+            observation_size=obs_size,
+            action_size=env.action_size,
+            normalize_observations=normalize_observations,
+            network_factory=network_factory,
+        )
         checkpoint.save(checkpoint_logdir, current_step, params, ckpt_config)
 
       # Run evals.

brax/training/checkpoint.py

@@ -71,6 +71,7 @@ def network_config(
       )
     del kwargs['preprocess_observations_fn']
   if 'activation' in kwargs:
+    # TODO: Add other activations.
     if kwargs['activation'] != defaults['activation']:
       raise ValueError('checkpointing only supports default activation')
     del kwargs['activation']

docs/release-notes/next-release.md

@@ -1 +1,3 @@
 # Brax Release Notes
+
+* Fix #595, patch for checkpointing with activations other than relu when no checkpoint path is specified.

pyproject.toml

@@ -71,6 +71,23 @@ Homepage = "http://github.com/google/brax"
 
 [tool.hatch.build.targets.wheel]
 packages = ["brax"]
+exclude = [
+  "/datasets",
+  "/docs",
+  "/notebooks",
+  "/tests",
+  "brax/experimental/barkour/assets",
+  "brax/experimental/barkour/data",
+]
+
+[tool.hatch.build.targets.sdist]
+exclude = [
+    "/datasets",
+    "/docs",
+    "/notebooks",
+    "brax/experimental/barkour/assets",
+    "brax/experimental/barkour/data",
+]
 
 [tool.isort]
 force_single_line = true