adding to device

scioip34 · May 13, 2024 · e65573b · e65573b
1 parent 1c01d39
commit e65573b
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/abm/app_madrl_foraging.py b/abm/app_madrl_foraging.py
@@ -93,7 +93,7 @@ def start(parallel=True, headless=False):
     if envconf["BRAIN_TYPE"] == "ideal":
         from abm.projects.madrl_foraging.madrl_simulation.heuristic_sims import HeuristicSimulation as Simulation
     else:
-        from abm.projects.madrl_foraging.madrl_simulation.madrl_sims_shared_replay import MADRLSimulation as Simulation
+        from abm.projects.madrl_foraging.madrl_simulation.madrl_sims import MADRLSimulation as Simulation
 
     vscreen_width = int(envconf["ENV_WIDTH"]) + 2 * int(envconf["WINDOW_PAD"]) + 10
     vscreen_height = int(envconf["ENV_HEIGHT"]) + 2 * int(envconf["WINDOW_PAD"]) + 10

diff --git a/abm/projects/madrl_foraging/madrl_simulation/madrl_sims.py b/abm/projects/madrl_foraging/madrl_simulation/madrl_sims.py
@@ -208,7 +208,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)
 
@@ -280,10 +280,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()