feat: configure checkpoints saving

scripts/hpc/train.pbs

@@ -67,7 +67,7 @@ fi
 python scripts/train.py \
     hydra.run.dir="$SCRATCH_RUNDIR" \
     ppo=stable \
-    logging=wandb_enabled
+    logging=hpc
 
 echo ">>> Staging out results to $DATA_RUNDIR..."
 cp -r "$SCRATCH_RUNDIR/." "$DATA_RUNDIR/"

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

scripts/train.py

@@ -41,6 +41,8 @@ def main(dict_cfg: DictConfig):
         entity=config.logging.wandb_entity,
         base_dir=os.path.dirname(run_dir),
         use_wandb=config.logging.track,
+        upload_final_model=config.logging.upload_final_model,
+        upload_checkpoints=config.logging.upload_checkpoints,
     )
     logger = get_logger()
     logger.info(f"Hydra-initialized run: {run_name}")