Fix adding in the wind bias

sisl · Mar 5, 2024 · f9e5982 · f9e5982
1 parent 820dcfa
commit f9e5982
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 1 deletion.
diff --git a/examples/training.gif b/examples/training.gif
diff --git a/pyrorl/pyrorl/envs/environment/environment.py b/pyrorl/pyrorl/envs/environment/environment.py
@@ -77,12 +77,13 @@ def __init__(
         self.evacuating_timestamps = np.full((num_rows, num_cols), np.inf)
 
         # Initialize placement of fire cells
+        self.num_fire_cells = num_fire_cells
         if custom_fire_locations:
             fire_rows = custom_fire_locations[:, 0]
             fire_cols = custom_fire_locations[:, 1]
             self.state_space[FIRE_INDEX, fire_rows, fire_cols] = 1
         else:
-            for _ in range(num_fire_cells):
+            for _ in range(self.num_fire_cells):
                 self.state_space[
                     FIRE_INDEX,
                     random.randint(0, num_rows - 1),

diff --git a/pyrorl/pyrorl/envs/pyrorl.py b/pyrorl/pyrorl/envs/pyrorl.py
@@ -40,6 +40,8 @@ def __init__(
         self.populated_areas = populated_areas
         self.paths = paths
         self.paths_to_pops = paths_to_pops
+        self.wind_speed = wind_speed
+        self.wind_angle = wind_angle
         self.fire_env = FireWorld(
             num_rows,
             num_cols,
@@ -78,6 +80,8 @@ def reset(self, seed=None, options=None):
             self.populated_areas,
             self.paths,
             self.paths_to_pops,
+            wind_speed = self.wind_speed,
+            wind_angle = self.wind_angle
         )
         state_space = self.fire_env.get_state()
         return state_space, {"": ""}