feat: configure checkpoints saving

scripts/simulate.py

@@ -57,6 +57,19 @@ def main(dict_cfg: DictConfig) -> None:
     nu = int(state.mj_model.nu)
 
     if model_path is not None:
+        # TODO: Refactoring Notice - The .flax checkpoint payload no longer encapsulates the config
+        # and no longer wraps parameters into a hardcoded list.
+        # The file now natively contains solely the pure raw Jax 'agent_state.params' FrozenDict mapping.
+        # The entire BrittleStarConfig is safely exported alongside it down at '..._metadata.yaml'.
+        #
+        # Example parsed layout from flax.serialization.from_bytes():
+        # {
+        #    'sensor_params': FrozenDict({...}),
+        #    'actor_params': FrozenDict({...}),
+        #    'critic_params': FrozenDict({...}),
+        #    ...
+        # }
+        # Update the RLModel.load function or subsequent destructuring to support this raw dictionary natively.
         policy = RLModel.load(Path(model_path))
         if hasattr(policy, "nu"):
             policy.nu = nu