""" Cat Rescue OpenEnv Environment ================================ A grid-based reinforcement learning environment where an AI agent navigates a grid to rescue trapped cats. Implements the OpenEnv standard interface: - reset() → initialise a fresh episode - step() → apply an action, get (observation, reward, done, info) - state() → return the current grid state Compatible with: - grader.py (reads observation dict + done/reward values) - rewards.py (can override REWARD_* constants or plug in its own function) - server.py (calls reset / step / state over HTTP) Grid cell values (also exported as constants so teammates can import them): EMPTY = 0 WALL = 1 CAT = 2 AGENT = 3 OBSTACLE = 4 (moving obstacle, Level 3 only) """ import copy import random from typing import Any, Dict, List, Optional, Tuple # --------------------------------------------------------------------------- # Try to import the openenv-core base class. # If the library is not installed yet we fall back to a plain object so the # file is still importable during development / local testing. # --------------------------------------------------------------------------- try: from openenv.core import BaseEnvironment # type: ignore except ImportError: class BaseEnvironment: # type: ignore """Minimal stand-in when openenv-core is not installed.""" pass # --------------------------------------------------------------------------- # Cell-type constants (import these in grader.py / rewards.py) # --------------------------------------------------------------------------- EMPTY = 0 WALL = 1 CAT = 2 AGENT = 3 OBSTACLE = 4 # --------------------------------------------------------------------------- # Action constants # --------------------------------------------------------------------------- UP = 0 DOWN = 1 LEFT = 2 RIGHT = 3 ACTION_NAMES = {UP: "UP", DOWN: "DOWN", LEFT: "LEFT", RIGHT: "RIGHT"} # --------------------------------------------------------------------------- # Reward constants (rewards.py can override these by monkey-patching or by # passing a custom rewards dict to CatRescueEnv.__init__) # --------------------------------------------------------------------------- REWARD_RESCUE_CAT = 1.0 # agent steps onto a cat cell REWARD_HIT_WALL = -0.3 # agent tries to walk into a wall / boundary REWARD_STEP = -0.1 # every step costs a little (encourages speed) REWARD_ALL_RESCUED = 2.0 # bonus when every cat on the level is saved # --------------------------------------------------------------------------- # Level definitions # Each level is a dict with: # grid_size : (rows, cols) # num_cats : how many cats are placed # walls : list of (row, col) wall positions (empty = no static walls) # num_moving : number of moving obstacles (Level 3 only) # --------------------------------------------------------------------------- LEVELS: Dict[int, Dict[str, Any]] = { 1: { "grid_size": (3, 3), "num_cats": 1, "walls": [], "num_moving": 0, }, 2: { "grid_size": (5, 5), "num_cats": 2, # Hand-crafted walls that create simple corridors "walls": [(1, 1), (1, 3), (3, 1), (3, 3)], "num_moving": 0, }, 3: { "grid_size": (7, 7), "num_cats": 3, # A more complex maze-like wall layout "walls": [ (1, 1), (1, 2), (1, 5), (2, 4), (3, 1), (3, 3), (3, 5), (4, 2), (5, 1), (5, 4), (5, 5), ], "num_moving": 2, }, } # --------------------------------------------------------------------------- # Main environment class # --------------------------------------------------------------------------- class CatRescueEnv(BaseEnvironment): """ Cat Rescue OpenEnv Environment. Parameters ---------- level : int Difficulty level (1, 2, or 3). max_steps : int Maximum number of steps per episode before the episode is forced to end (prevents infinite loops). seed : int | None Random seed for reproducibility. custom_rewards : dict | None Optional dict to override default reward values. Keys: 'rescue_cat', 'hit_wall', 'step', 'all_rescued' """ # ------------------------------------------------------------------ # # Constructor # # ------------------------------------------------------------------ # def __init__( self, level: int = 1, max_steps: int = 200, seed: Optional[int] = None, custom_rewards: Optional[Dict[str, float]] = None, ): if level not in LEVELS: raise ValueError(f"Level must be 1, 2, or 3. Got {level}.") self.level = level self.max_steps = max_steps self._rng = random.Random(seed) # Load level config cfg = LEVELS[level] self.rows, self.cols = cfg["grid_size"] self._static_walls = set(map(tuple, cfg["walls"])) self._num_cats = cfg["num_cats"] self._num_moving = cfg["num_moving"] # Allow rewards.py to supply its own reward schedule self._rewards = { "rescue_cat": REWARD_RESCUE_CAT, "hit_wall": REWARD_HIT_WALL, "step": REWARD_STEP, "all_rescued": REWARD_ALL_RESCUED, } if custom_rewards: self._rewards.update(custom_rewards) # Internal state (populated by reset()) self.grid: List[List[int]] = [] self.agent_pos: Tuple[int, int] = (0, 0) self.cat_positions: List[Tuple[int, int]] = [] self.obstacle_positions: List[Tuple[int, int]] = [] self._obstacle_dirs: List[Tuple[int, int]] = [] # velocity per obstacle self.cats_rescued: int = 0 self.step_count: int = 0 self.done: bool = False # Run an initial reset so the env is ready to use immediately self.reset() # ------------------------------------------------------------------ # # OpenEnv Interface — reset() # # ------------------------------------------------------------------ # def reset(self) -> Dict[str, Any]: """ Reset the environment to a fresh starting state. Places the agent at (0, 0), randomly distributes cats on empty cells, places static walls (Level 2+), and spawns moving obstacles (Level 3). Returns ------- observation : dict The initial observation (same format as step() returns). """ # 1. Build a blank grid filled with EMPTY cells self.grid = [[EMPTY] * self.cols for _ in range(self.rows)] # 2. Place static walls (Level 2 and 3) for (r, c) in self._static_walls: self.grid[r][c] = WALL # 3. Collect all free cells (not walls) for random placement free_cells = self._get_free_cells() # 4. Place the agent at (0, 0) — always the top-left corner self.agent_pos = (0, 0) self.grid[0][0] = AGENT free_cells.discard((0, 0)) # 5. Place cats randomly on remaining free cells self.cat_positions = [] for _ in range(self._num_cats): if not free_cells: break # safety guard — shouldn't happen with valid configs pos = self._rng.choice(list(free_cells)) free_cells.discard(pos) self.cat_positions.append(pos) self.grid[pos[0]][pos[1]] = CAT # 6. Place moving obstacles (Level 3 only) self.obstacle_positions = [] self._obstacle_dirs = [] directions = [(0, 1), (0, -1), (1, 0), (-1, 0)] for _ in range(self._num_moving): if not free_cells: break pos = self._rng.choice(list(free_cells)) free_cells.discard(pos) self.obstacle_positions.append(list(pos)) # mutable list self._obstacle_dirs.append(self._rng.choice(directions)) self.grid[pos[0]][pos[1]] = OBSTACLE # 7. Reset counters self.cats_rescued = 0 self.step_count = 0 self.done = False return self.state() # ------------------------------------------------------------------ # # OpenEnv Interface — step() # # ------------------------------------------------------------------ # def step(self, action: int) -> Tuple[Dict[str, Any], float, bool, Dict[str, Any]]: """ Execute one action in the environment. Parameters ---------- action : int One of UP (0), DOWN (1), LEFT (2), RIGHT (3). Returns ------- observation : dict Current state of the environment (from state()). reward : float Reward received for this step. done : bool True when the episode has ended (all cats rescued or max_steps hit). info : dict Extra diagnostic data useful for grader.py and debugging. Keys: 'cats_rescued', 'step_count', 'agent_pos', 'action_name', 'event' """ if self.done: # Episode already finished — return current state with no reward return self.state(), 0.0, self.done, self._make_info("episode_done") reward = self._rewards["step"] # base cost for every step event = "step" # human-readable event label # --- Compute the target cell ----------------------------------- row, col = self.agent_pos d_row, d_col = self._action_delta(action) new_row = row + d_row new_col = col + d_col # --- Check if the move is valid -------------------------------- if self._is_out_of_bounds(new_row, new_col) or self._is_wall(new_row, new_col): # Agent bumps into boundary or wall — penalise and stay put reward += self._rewards["hit_wall"] event = "hit_wall" else: # Valid move — update agent position on the grid cell_content = self.grid[new_row][new_col] # Clear the agent's old position self.grid[row][col] = EMPTY # Check what is in the destination cell if cell_content == CAT: # Cat rescued! self.cats_rescued += 1 self.cat_positions.remove((new_row, new_col)) reward += self._rewards["rescue_cat"] event = "cat_rescued" # Check if ALL cats have been rescued if self.cats_rescued == self._num_cats: reward += self._rewards["all_rescued"] self.done = True event = "all_rescued" elif cell_content == OBSTACLE: # Ran into a moving obstacle — same penalty as a wall reward += self._rewards["hit_wall"] event = "hit_obstacle" # Agent stays at current position (don't move into obstacle) self.grid[row][col] = AGENT self.step_count += 1 return ( self.state(), reward, self.done, self._make_info(event, action), ) # Move agent to new position self.agent_pos = (new_row, new_col) self.grid[new_row][new_col] = AGENT # --- Move obstacles (Level 3) ---------------------------------- if self._num_moving > 0: self._move_obstacles() # --- Check step limit ------------------------------------------ self.step_count += 1 if self.step_count >= self.max_steps and not self.done: self.done = True event = "max_steps_reached" info = self._make_info(event, action) return self.state(), reward, self.done, info # ------------------------------------------------------------------ # # OpenEnv Interface — state() # # ------------------------------------------------------------------ # def state(self) -> Dict[str, Any]: """ Return the current environment state as an observation dict. This is what grader.py reads to evaluate the agent. Returns ------- dict with keys: 'grid' : 2-D list of int (deep copy — safe to mutate) 'agent_pos' : (row, col) tuple 'cat_positions' : list of (row, col) tuples 'cats_rescued' : int — how many cats saved so far 'cats_remaining' : int — how many cats still on the grid 'step_count' : int 'done' : bool 'level' : int 'grid_size' : (rows, cols) """ return { "grid": copy.deepcopy(self.grid), "agent_pos": self.agent_pos, "cat_positions": list(self.cat_positions), "cats_rescued": self.cats_rescued, "cats_remaining": self._num_cats - self.cats_rescued, "step_count": self.step_count, "done": self.done, "level": self.level, "grid_size": (self.rows, self.cols), } # ------------------------------------------------------------------ # # Private helpers # # ------------------------------------------------------------------ # def _action_delta(self, action: int) -> Tuple[int, int]: """Map an action integer to a (row_delta, col_delta) pair.""" return { UP: (-1, 0), DOWN: ( 1, 0), LEFT: ( 0, -1), RIGHT: ( 0, 1), }.get(action, (0, 0)) # (0, 0) → no-op for unknown actions def _is_out_of_bounds(self, row: int, col: int) -> bool: """Return True if (row, col) is outside the grid.""" return not (0 <= row < self.rows and 0 <= col < self.cols) def _is_wall(self, row: int, col: int) -> bool: """Return True if the cell at (row, col) contains a static wall.""" return self.grid[row][col] == WALL def _get_free_cells(self) -> set: """Return a set of (row, col) positions that are currently EMPTY.""" free = set() for r in range(self.rows): for c in range(self.cols): if self.grid[r][c] == EMPTY: free.add((r, c)) return free def _move_obstacles(self) -> None: """ Move each obstacle one step in its current direction. If the next cell is a wall, boundary, or another obstacle the obstacle bounces — its direction is reversed. The agent is never overwritten: obstacles bounce off the agent cell too, keeping collision detection clean. """ for i, (obs_pos, obs_dir) in enumerate( zip(self.obstacle_positions, self._obstacle_dirs) ): r, c = obs_pos dr, dc = obs_dir nr, nc = r + dr, c + dc # Clear current obstacle cell (only if it hasn't been overwritten # by the agent moving there in the same step) if self.grid[r][c] == OBSTACLE: self.grid[r][c] = EMPTY # Decide whether to bounce or advance blocked = ( self._is_out_of_bounds(nr, nc) or self.grid[nr][nc] in (WALL, OBSTACLE, AGENT) ) if blocked: # Reverse direction (bounce) self._obstacle_dirs[i] = (-dr, -dc) nr, nc = r, c # stay in place this tick self.obstacle_positions[i] = [nr, nc] self.grid[nr][nc] = OBSTACLE def _make_info(self, event: str, action: Optional[int] = None) -> Dict[str, Any]: """Build the info dict returned alongside every step.""" return { "cats_rescued": self.cats_rescued, "step_count": self.step_count, "agent_pos": self.agent_pos, "action_name": ACTION_NAMES.get(action, "N/A") if action is not None else "N/A", "event": event, } # ------------------------------------------------------------------ # # Convenience / debug helpers (not part of the OpenEnv interface) # # ------------------------------------------------------------------ # def render(self) -> str: """ Return a human-readable ASCII string of the current grid. Useful for debugging and demos. Symbols: . EMPTY # WALL C CAT A AGENT X MOVING OBSTACLE """ symbols = {EMPTY: ".", WALL: "#", CAT: "C", AGENT: "A", OBSTACLE: "X"} lines = [] for row in self.grid: lines.append(" ".join(symbols.get(cell, "?") for cell in row)) header = f"Level {self.level} | Step {self.step_count} | Cats rescued {self.cats_rescued}/{self._num_cats}" return header + "\n" + "\n".join(lines) def action_space(self) -> List[int]: """Return the list of valid action integers.""" return [UP, DOWN, LEFT, RIGHT] def observation_space(self) -> Dict[str, Any]: """ Describe the observation space for grader.py / server.py. Returns a metadata dict (not a sampled observation). """ return { "grid_shape": (self.rows, self.cols), "cell_types": {"EMPTY": EMPTY, "WALL": WALL, "CAT": CAT, "AGENT": AGENT, "OBSTACLE": OBSTACLE}, "agent_pos": "Tuple[int, int]", "cat_positions": "List[Tuple[int, int]]", "cats_rescued": "int", "step_count": "int", "done": "bool", } # --------------------------------------------------------------------------- # Quick smoke-test — run `python environment.py` to verify everything works # --------------------------------------------------------------------------- if __name__ == "__main__": for lvl in (1, 2, 3): print(f"\n{'='*40}") print(f" Smoke-test — Level {lvl}") print(f"{'='*40}") env = CatRescueEnv(level=lvl, seed=42) print(env.render()) print() # Take 10 random steps for _ in range(10): action = random.choice(env.action_space()) obs, reward, done, info = env.step(action) print(f"Action: {ACTION_NAMES[action]:<5} | " f"Reward: {reward:+.1f} | " f"Event: {info['event']:<16} | " f"Agent: {info['agent_pos']} | " f"Rescued: {info['cats_rescued']}") if done: print(" >> Episode finished early!") break print() print(env.render())