"putting back shared experience"

scioip34 · May 13, 2024 · 2991c47 · 2991c47
1 parent d6d9122
commit 2991c47
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/abm/projects/madrl_foraging/madrl_simulation/madrl_sims_shared_replay.py b/abm/projects/madrl_foraging/madrl_simulation/madrl_sims_shared_replay.py
@@ -207,7 +207,7 @@ def initialize_environment(self):
                 # calculate number of resources left in the patch
                 resc = list(ag_resc_overlap.keys())[0]
                 ag.policy_network.state_tensor = torch.FloatTensor(
-                ag.soc_v_field.tolist() + [resc.resc_left / resc.resc_units]).unsqueeze(0)
+                ag.soc_v_field.tolist() + [resc.resc_left / resc.resc_units]).unsqueeze(0).to(device)
             else:
                 ag.policy_network.state_tensor = torch.FloatTensor(ag.soc_v_field.tolist() + [0.0]).unsqueeze(0).to(device)
 
@@ -279,10 +279,10 @@ def start_madqn(self):
                             ag_resc_overlap = self.agent_resource_overlap([ag])
                             resc= list(ag_resc_overlap.keys())[0]
 
-                            ag.policy_network.next_state_tensor = torch.FloatTensor(ag.soc_v_field.tolist() + [resc.resc_left/resc.resc_units]).unsqueeze(0)
+                            ag.policy_network.next_state_tensor = torch.FloatTensor(ag.soc_v_field.tolist() + [resc.resc_left/resc.resc_units]).unsqueeze(0).to(device)
 
                         else:
-                            ag.policy_network.next_state_tensor = torch.FloatTensor(ag.soc_v_field.tolist() + [0.0]).unsqueeze(0)
+                            ag.policy_network.next_state_tensor = torch.FloatTensor(ag.soc_v_field.tolist() + [0.0]).unsqueeze(0).to(device)
 
                         # Calculate the reward as a weighted sum of the individual and collective search efficiency
                         reward = ag.compute_reward()