Add model loading for other algorithms

UoA-CARES · Sep 5, 2023 · 88f697f · 88f697f
1 parent a1a1035
commit 88f697f
Showing 1 changed file with 21 additions and 0 deletions.
diff --git a/cares_reinforcement_learning/util/NetworkFactory.py b/cares_reinforcement_learning/util/NetworkFactory.py
@@ -7,6 +7,9 @@ def create_DQN(args):
 
     network = Network(args["observation_size"], args["action_num"], args["lr"])
 
+    if "network_file_path" in args:
+        network.load_state_dict(torch.load(args["network_file_path"]))
+
     agent = DQN(
         network=network,
         gamma=args["gamma"],
@@ -21,6 +24,9 @@ def create_DuelingDQN(args):
 
     network = DuelingNetwork(args["observation_size"], args["action_num"], args["lr"])
 
+    if "network_file_path" in args:
+        network.load_state_dict(torch.load(args["network_file_path"]))
+
     agent = DQN(
         network=network,
         gamma=args["gamma"],
@@ -35,6 +41,9 @@ def create_DDQN(args):
 
     network = Network(args["observation_size"], args["action_num"], args["lr"])
 
+    if "network_file_path" in args:
+        network.load_state_dict(torch.load(args["network_file_path"]))
+
     agent = DoubleDQN(
         network=network,
         gamma=args["gamma"],
@@ -52,6 +61,10 @@ def create_PPO(args):
     actor = Actor(args["observation_size"], args["action_num"], args["actor_lr"])
     critic = Critic(args["observation_size"], args["critic_lr"])
 
+    if "actor_file_path" in args and "critic_file_path" in args:
+        actor.load_state_dict(torch.load(args["actor_file_path"]))
+        critic.load_state_dict(torch.load(args["critic_file_path"]))
+
     agent = PPO(
         actor_network=actor,
         critic_network=critic,
@@ -70,6 +83,10 @@ def create_SAC(args):
     actor = Actor(args["observation_size"], args["action_num"], args["actor_lr"])
     critic = Critic(args["observation_size"], args["action_num"], args["critic_lr"])
 
+    if "actor_file_path" in args and "critic_file_path" in args:
+        actor.load_state_dict(torch.load(args["actor_file_path"]))
+        critic.load_state_dict(torch.load(args["critic_file_path"]))
+
     agent = SAC(
         actor_network=actor,
         critic_network=critic,
@@ -89,6 +106,10 @@ def create_DDPG(args):
     actor = Actor(args["observation_size"], args["action_num"], args["actor_lr"])
     critic = Critic(args["observation_size"], args["action_num"], args["critic_lr"])
 
+    if "actor_file_path" in args and "critic_file_path" in args:
+        actor.load_state_dict(torch.load(args["actor_file_path"]))
+        critic.load_state_dict(torch.load(args["critic_file_path"]))
+
     agent = DDPG(
         actor_network=actor,
         critic_network=critic,