Update sac_agent.py

kelvin34501 · Jun 27, 2022 · a9336bd · a9336bd
1 parent d5290c6
commit a9336bd
Showing 1 changed file with 3 additions and 4 deletions.
diff --git a/rl_games/algos_torch/sac_agent.py b/rl_games/algos_torch/sac_agent.py
@@ -196,9 +196,8 @@ def get_weights(self):
         return state
 
     def save(self, fn):
-        pass
-        #state = self.get_full_state_weights()
-        #torch_ext.save_checkpoint(fn, state)
+        state = self.get_full_state_weights()
+        torch_ext.save_checkpoint(fn, state)
 
     def set_weights(self, weights):
         self.model.sac_network.actor.load_state_dict(weights['actor'])
@@ -544,4 +543,4 @@ def train(self):
                     print('MAX EPOCHS NUM!')
                     return self.last_mean_rewards, self.epoch_num
 
-                update_time = 0
+                update_time = 0