The improved model cannot load properly (#475)

* Modify AgentBase, config, and run scripts

* Modify AgentBase, config, and run scripts

---------

Co-authored-by: ty <[email protected]>

Author: kernel-tye

elegantrl/agents/AgentBase.py

@@ -292,7 +292,7 @@ def save_or_load_agent(self, cwd: str, if_save: bool):
                 continue
 
             if if_save:
-                th.save(getattr(self, attr_name).state_dict(), file_path)
+                th.save(getattr(self, attr_name), file_path)
             elif os.path.isfile(file_path):
                 setattr(self, attr_name, th.load(file_path, map_location=self.device))
 

elegantrl/train/config.py

@@ -42,6 +42,7 @@ def __init__(self, agent_class=None, env_class=None, env_args=None):
         self.clip_grad_norm = 3.0  # 0.1 ~ 4.0, clip the gradient after normalization
         self.state_value_tau = 0  # the tau of normalize for value and state `std = (1-std)*std + tau*std`
         self.soft_update_tau = 5e-3  # 2 ** -8 ~= 5e-3. the tau of soft target update `net = (1-tau)*net + tau*net1`
+        self.continue_train = False  # continue train use last train saved models
         if self.if_off_policy:  # off-policy
             self.batch_size = int(64)  # num of transitions sampled from replay buffer.
             self.horizon_len = int(512)  # collect horizon_len step while exploring, then update networks

elegantrl/train/run.py

@@ -45,7 +45,8 @@ def train_agent_single_process(args: Config):
 
     '''init agent'''
     agent = args.agent_class(args.net_dims, args.state_dim, args.action_dim, gpu_id=args.gpu_id, args=args)
-    agent.save_or_load_agent(args.cwd, if_save=False)
+    if args.continue_train:
+        agent.save_or_load_agent(args.cwd, if_save=False)
 
     '''init agent.last_state'''
     state, info_dict = env.reset()
@@ -234,7 +235,8 @@ def run(self):
 
         '''Learner init agent'''
         agent = args.agent_class(args.net_dims, args.state_dim, args.action_dim, gpu_id=args.gpu_id, args=args)
-        agent.save_or_load_agent(args.cwd, if_save=False)
+        if args.continue_train:
+            agent.save_or_load_agent(args.cwd, if_save=False)
 
         '''Learner init buffer'''
         if args.if_off_policy:
@@ -373,7 +375,8 @@ def run(self):
 
         '''init agent'''
         agent = args.agent_class(args.net_dims, args.state_dim, args.action_dim, gpu_id=args.gpu_id, args=args)
-        agent.save_or_load_agent(args.cwd, if_save=False)
+        if args.continue_train:
+            agent.save_or_load_agent(args.cwd, if_save=False)
 
         '''init agent.last_state'''
         state, info_dict = env.reset()