fix: fix compensator saving

Author: Nips20262

omnisafe/algorithms/off_policy/ddpg_cbf.py

@@ -51,7 +51,7 @@ def _init_env(self) -> None:
             self._seed,
             self._cfgs,
         )
-        solver = PendulumSolver(device=self._cfgs.train_cfgs.device)
+        solver = PendulumSolver(device=self._device)
         compensator = BarrierCompensator(
             obs_dim=self._env.observation_space.shape[0],
             act_dim=self._env.action_space.shape[0],
@@ -120,11 +120,18 @@ def _specific_save(self) -> None:
             os.makedirs(os.path.dirname(path), exist_ok=True)
             joblib.dump(self._env.gp_models, path)
 
-    def _log_what_to_save(self) -> dict[str, Any]:
-        """Define what need to be saved below."""
+    def _setup_torch_saver(self) -> None:
+        """Define what need to be saved below.
+
+        OmniSafe's main storage interface is based on PyTorch. If you need to save models in other
+        formats, please use :meth:`_specific_save`.
+        """
         what_to_save: dict[str, Any] = {}
 
         what_to_save['pi'] = self._actor_critic.actor
         what_to_save['compensator'] = self._env.compensator
+        if self._cfgs.algo_cfgs.obs_normalize:
+            obs_normalizer = self._env.save()['obs_normalizer']
+            what_to_save['obs_normalizer'] = obs_normalizer
 
         self._logger.setup_torch_saver(what_to_save)