Upload folder using huggingface_hub

env/firefighter_env.py

@@ -0,0 +1,102 @@
+# firefighter_env.py
+import copy
+import gymnasium as gym
+from gymnasium import spaces
+import numpy as np
+
+class FireFighterEnv(gym.Env):
+    metadata = {"render_modes": ["human"], "render_fps": 4}
+
+    def __init__(self, render_mode=None):
+        super(FireFighterEnv, self).__init__()
+
+        self.grid_size = 4
+        self.max_steps = 60
+
+        self.agent_start = (0, 0)
+        self.bucket_pos = (1, 1)
+        self.fire_pos = (1, 3)
+        self.goal_pos = (3, 3)
+        self.walls = {(1, 2), (2, 1)}
+
+        # Actions: 0=Up, 1=Down, 2=Left, 3=Right
+        self.action_space = spaces.Discrete(4)
+
+        # Observation: (x, y, has_bucket, fire_out)
+        self.observation_space = spaces.MultiDiscrete([4, 4, 2, 2])
+
+        self.render_mode = render_mode
+        self.reset()
+
+    def reset(self, seed=None, options=None):
+        super().reset(seed=seed)
+
+        self.agent_pos = list(self.agent_start)
+        self.has_bucket = False
+        self.fire_out = False
+        self.steps = 0
+
+        obs = self._get_obs()
+        return obs, {}
+
+    def step(self, action):
+        self.steps += 1
+        x, y = self.agent_pos
+
+        move = {0: (-1, 0), 1: (1, 0), 2: (0, -1), 3: (0, 1)}
+        dx, dy = move[action]
+
+        # Stochastic transitions
+        if np.random.rand() > 0.8:
+            dx, dy = move[np.random.choice([a for a in move if a != action])]
+
+        new_x = np.clip(x + dx, 0, self.grid_size - 1)
+        new_y = np.clip(y + dy, 0, self.grid_size - 1)
+
+        if (new_x, new_y) in self.walls:
+            reward = -5
+        else:
+            self.agent_pos = [new_x, new_y]
+            reward = 0
+
+        # Bucket collection
+        if tuple(self.agent_pos) == self.bucket_pos and not self.has_bucket:
+            self.has_bucket = True
+            reward += 10
+
+        # Extinguish fire
+        if tuple(self.agent_pos) == self.fire_pos and self.has_bucket and not self.fire_out:
+            self.fire_out = True
+            reward += 10
+
+        # Reaching goal
+        if tuple(self.agent_pos) == self.goal_pos:
+            if self.fire_out:
+                reward += 10
+                terminated = True
+            else:
+                reward -= 10
+                terminated = True
+        else:
+            terminated = False
+
+        truncated = self.steps >= self.max_steps
+
+        obs = self._get_obs()
+        info = {}
+        return obs, reward, terminated, truncated, info
+
+    def _get_obs(self):
+        return np.array([*self.agent_pos, int(self.has_bucket), int(self.fire_out)], dtype=np.int32)
+    def save_state(self):
+        return copy.deepcopy(self)
+
+    def reset_to(self, saved_env):
+        self.__dict__.update(saved_env.__dict__)
+
+
+    def render(self):
+        pass
+
+    def close(self):
+        pass                

env/renderer.py

@@ -0,0 +1,106 @@
+import matplotlib
+matplotlib.use('Agg')
+import matplotlib.pyplot as plt
+import matplotlib.image as mpimg
+import numpy as np
+import os
+
+class FireFighterRenderer:
+    def __init__(self, save_dir=None):
+        self.grid_size = 4
+        self.save_dir = save_dir
+        self.frame_idx = 0
+        self.sprite_dir = "assets/sprites"
+        self.rewards = []
+
+        self.fig, (self.ax_grid, self.ax_plot) = plt.subplots(1, 2, figsize=(10, 5),
+                                                              gridspec_kw={'width_ratios': [1, 1]})
+
+        # Load sprites
+        def load(name):
+            path = os.path.join(self.sprite_dir, f"{name}.png")
+            img = mpimg.imread(path)
+            assert img is not None, f"Failed to load: {path}"
+            return img
+
+        self.sprites = {
+            "robot_white": load("robot_white"),
+            "robot_blue": load("robot_blue"),
+            "robot_green": load("robot_green"),
+            "bucket": load("bucket"),
+            "fire": load("fire"),
+            "goal": load("goal"),
+            "wall": load("wall"),
+        }
+
+        if save_dir:
+            os.makedirs(save_dir, exist_ok=True)
+
+    def render(self, agent_pos, has_bucket, fire_out, bucket_pos, fire_pos, goal_pos, walls, reward):
+        self.ax_grid.clear()
+        self.ax_plot.clear()
+        self.rewards.append(reward)
+
+        self.ax_grid.set_xlim(0, self.grid_size)
+        self.ax_grid.set_ylim(0, self.grid_size)
+        self.ax_grid.set_xticks([])
+        self.ax_grid.set_yticks([])
+        self.ax_grid.set_aspect('equal')
+
+        # White background grid
+        for x in range(self.grid_size):
+            for y in range(self.grid_size):
+                self.ax_grid.add_patch(plt.Rectangle((y, self.grid_size - 1 - x), 1, 1,
+                                                     edgecolor='black', facecolor='white', linewidth=1))
+
+        # Function to place sprites
+        def draw(sprite_name, x, y):
+            sprite = self.sprites[sprite_name]
+            self.ax_grid.imshow(sprite,
+                                extent=(y, y + 1, self.grid_size - 1 - x, self.grid_size - x),
+                                zorder=10)
+
+        # Draw elements
+        for wx, wy in walls:
+            draw("wall", wx, wy)
+
+        if not fire_out:
+            fx, fy = fire_pos
+            draw("fire", fx, fy)
+
+        bx, by = bucket_pos
+        draw("bucket", bx, by)
+
+        gx, gy = goal_pos
+        draw("goal", gx, gy)
+
+        ax, ay = agent_pos
+        if has_bucket and not fire_out:
+            robot_color = "robot_blue"
+        elif fire_out:
+            robot_color = "robot_green"
+        else:
+            robot_color = "robot_white"
+        draw(robot_color, ax, ay)
+
+        self.ax_grid.set_title(f"Step {self.frame_idx}")
+
+        # Reward plot
+        self.ax_plot.plot(np.cumsum(self.rewards), color='green', marker='o')
+        self.ax_plot.set_title("Cumulative Reward")
+        self.ax_plot.set_xlabel("Step")
+        self.ax_plot.set_ylabel("Total Reward")
+        self.ax_plot.grid(True)
+
+        if self.save_dir:
+            frame_path = os.path.join(self.save_dir, f"frame_{self.frame_idx:03d}.png")
+            self.fig.tight_layout()
+            self.fig.savefig(frame_path)
+            self.frame_idx += 1
+        else:
+            plt.pause(0.3)
+            plt.draw()
+
+    def close(self):
+        plt.close(self.fig)
+

env/utils.py

@@ -0,0 +1,10 @@
+def format_observation(obs):
+    return {
+        'x': obs[0],
+        'y': obs[1],
+        'has_bucket': bool(obs[2]),
+        'fire_out': bool(obs[3]),
+    }
+
+def is_terminal_state(obs):
+    return (obs[0], obs[1]) == (3, 3) and obs[3] == 1