reward classifier fix checkpoint save load

Signed-off-by: youliang <[email protected]>

examples/async_bin_relocation_fwbw_drq/train_reward_classifier.py

@@ -1,6 +1,7 @@
 import pickle as pkl
 import jax
 from jax import numpy as jnp
+import flax
 import flax.linen as nn
 from flax.training import checkpoints
 import optax
@@ -161,6 +162,8 @@ def loss_fn(params):
             f"Epoch: {epoch+1}, Train Loss: {train_loss:.4f}, Train Accuracy: {train_accuracy:.4f}"
         )
 
+    # this is used to save the without the orbax checkpointing
+    flax.config.update('flax_use_orbax_checkpointing', False)
     checkpoints.save_checkpoint(
         FLAGS.classifier_ckpt_path,
         classifier,

examples/async_cable_route_drq/train_reward_classifier.py

@@ -1,6 +1,7 @@
 import pickle as pkl
 import jax
 from jax import numpy as jnp
+import flax
 import flax.linen as nn
 from flax.training import checkpoints
 import optax
@@ -161,7 +162,8 @@ def loss_fn(params):
         print(
             f"Epoch: {epoch+1}, Train Loss: {train_loss:.4f}, Train Accuracy: {train_accuracy:.4f}"
         )
-
+    # this is used to save the without the orbax checkpointing
+    flax.config.update('flax_use_orbax_checkpointing', False)
     checkpoints.save_checkpoint(
         FLAGS.classifier_ckpt_path,
         classifier,

serl_launcher/serl_launcher/networks/reward_classifier.py

@@ -5,7 +5,7 @@
 from flax.training.train_state import TrainState
 from flax.training import checkpoints
 import optax
-from typing import Callable, Dict, List
+from typing import Callable, Dict, List, Optional
 
 
 from serl_launcher.vision.resnet_v1 import resnetv1_configs, PreTrainedResNetEncoder
@@ -94,6 +94,7 @@ def load_classifier_func(
     sample: Dict,
     image_keys: List[str],
     checkpoint_path: str,
+    step: Optional[int] = None
 ) -> Callable[[Dict], jnp.ndarray]:
     """
     Return: a function that takes in an observation
@@ -103,7 +104,7 @@ def load_classifier_func(
     classifier = checkpoints.restore_checkpoint(
         checkpoint_path,
         target=classifier,
-        step=100,
+        step=step,
     )
     func = lambda obs: classifier.apply_fn(
         {"params": classifier.params}, obs, train=False