from typing import Dict, List, Tuple from collections import deque import copy from models import Observation, Action, EpisodeResult, PlacedFurniture, Furniture from tasks import TASKS NEAR_WINDOW_RADIUS = 2 NEAR_DOOR_RADIUS = 3 class InteriorEnv: def __init__(self): self.task_id = "" self.room_size = (0, 0) self.door_position = (0, 0) self.windows: List[Tuple[int, int]] = [] self.item_defs: Dict[str, Furniture] = {} self.unplaced_items: Dict[str, Furniture] = {} self.placed_items: Dict[str, PlacedFurniture] = {} self.max_steps = 10 self.step_count = 0 self.success = False self.logs: List[str] = [] self.solar_grid: List[List[bool]] = [] self.thermal_grid: List[List[float]] = [] def reset(self, task_config: dict, task_id: str) -> Observation: self.task_id = task_id self.room_size = task_config["room_size"] self.door_position = task_config["door_position"] self.windows = task_config.get("windows", []) self.item_defs = {item.id: item for item in task_config["items"]} self.unplaced_items = {item.id: item for item in task_config["items"]} self.placed_items = {} self.max_steps = task_config["max_steps"] self.step_count = 0 self.success = False self.solar_grid = [[False] * self.room_size[0] for _ in range(self.room_size[1])] self.thermal_grid = [[0.0] * self.room_size[0] for _ in range(self.room_size[1])] self.logs = self._build_start_logs(task_config) return self._get_obs() def _build_start_logs(self, cfg: dict) -> List[str]: lines = [ f"START: Task={self.task_id.upper()} | Room={self.room_size} | Door={self.door_position}", f"Windows at: {self.windows}", f"Items to place: {len(self.unplaced_items)}", "─── Design Rules ────────────────────────────────", ] for item in cfg["items"]: for rule in item.rules: lines.append(f" [{item.id.upper()}]: {self._describe_rule(rule)}") lines.append("─────────────────────────────────────────────────") return lines def _describe_rule(self, rule: str) -> str: if rule == "must_touch_wall": return "Must be placed flush against any wall." if rule == "near_window": return f"Must be within {NEAR_WINDOW_RADIUS} tiles of a window." if rule == "all_sides_clear": return "Must have at least 1 free tile on ALL 4 sides (e.g., dining space)." if rule == "near_door": return f"Must be within {NEAR_DOOR_RADIUS} tiles of the main entry door." if rule.startswith("adjacent_to:"): target = rule.split(":")[1] return f"Must be placed directly adjacent (touching) the {target.upper()}." if rule.startswith("must_face:"): target = rule.split(":")[1] return f"Must align on the same axis as the {target.upper()} (facing it)." return rule # ─── Geometry ───────────────────────────────────────────────────────────── def _bbox(self, placed: PlacedFurniture, item_def: Furniture) -> Tuple[int, int, int, int]: w = item_def.width if placed.rotation == 0 else item_def.height h = item_def.height if placed.rotation == 0 else item_def.width return (placed.x, placed.x + w - 1, placed.y, placed.y + h - 1) def _bbox_action(self, action: Action, item_def: Furniture) -> Tuple[int, int, int, int]: w = item_def.width if action.rotation == 0 else item_def.height h = item_def.height if action.rotation == 0 else item_def.width return (action.x, action.x + w - 1, action.y, action.y + h - 1) def _build_grid(self) -> List[List[bool]]: rw, rh = self.room_size grid = [[False] * rw for _ in range(rh)] for iid, placed in self.placed_items.items(): xb, xe, yb, ye = self._bbox(placed, self.item_defs[iid]) for gy in range(yb, ye + 1): for gx in range(xb, xe + 1): if 0 <= gx < rw and 0 <= gy < rh: grid[gy][gx] = True return grid def _bfs_reachable(self, grid: List[List[bool]]) -> List[List[bool]]: rw, rh = self.room_size visited = [[False] * rw for _ in range(rh)] dx, dy = self.door_position if grid[dy][dx]: return visited queue = deque([(dx, dy)]) visited[dy][dx] = True while queue: cx, cy = queue.popleft() for nx, ny in [(cx+1,cy),(cx-1,cy),(cx,cy+1),(cx,cy-1)]: if 0 <= nx < rw and 0 <= ny < rh and not grid[ny][nx] and not visited[ny][nx]: visited[ny][nx] = True queue.append((nx, ny)) return visited def _all_neighbors(self, pid: str) -> List[Tuple[int, int]]: placed = self.placed_items[pid] xb, xe, yb, ye = self._bbox(placed, self.item_defs[pid]) rw, rh = self.room_size result = [] for gy in range(yb, ye + 1): for gx in range(xb, xe + 1): for nx, ny in [(gx+1,gy),(gx-1,gy),(gx,gy+1),(gx,gy-1)]: if 0 <= nx < rw and 0 <= ny < rh: result.append((nx, ny)) return result def _build_solar_grid(self) -> List[List[bool]]: rw, rh = self.room_size solar = [[False] * rw for _ in range(rh)] tall_grid = [[False] * rw for _ in range(rh)] for iid, placed in self.placed_items.items(): if any(r == "is_tall" for r in self.item_defs[iid].rules): xb, xe, yb, ye = self._bbox(placed, self.item_defs[iid]) for gy in range(yb, ye + 1): for gx in range(xb, xe + 1): if 0 <= gx < rw and 0 <= gy < rh: tall_grid[gy][gx] = True for wx, wy in self.windows: for dx, dy in [(1,0), (-1,0), (0,1), (0,-1)]: cx, cy = wx, wy while 0 <= cx < rw and 0 <= cy < rh: if tall_grid[cy][cx]: break solar[cy][cx] = True cx += dx cy += dy return solar def _build_thermal_grid(self) -> List[List[float]]: rw, rh = self.room_size thermal = [[0.0] * rw for _ in range(rh)] cool_sources = self.windows + [self.door_position] for cx, cy in cool_sources: for gy in range(rh): for gx in range(rw): dist = abs(gx - cx) + abs(gy - cy) if dist <= 3: thermal[gy][gx] -= 2.0 / (dist + 1) for iid, placed in self.placed_items.items(): for rule in self.item_defs[iid].rules: if rule.startswith("emits_heat:") or rule.startswith("emits_cold:"): val = float(rule.split(":")[1]) if rule.startswith("emits_cold:"): val = -val xb, xe, yb, ye = self._bbox(placed, self.item_defs[iid]) for gy in range(rh): for gx in range(rw): dx = max(0, xb - gx, gx - xe) dy = max(0, yb - gy, gy - ye) dist = dx + dy if dist <= 3: thermal[gy][gx] += val / (dist + 1) return thermal # ─── Rule Checkers ──────────────────────────────────────────────────────── def _check_must_touch_wall(self, pid: str) -> bool: rw, rh = self.room_size placed = self.placed_items[pid] xb, xe, yb, ye = self._bbox(placed, self.item_defs[pid]) return xb == 0 or xe == rw - 1 or yb == 0 or ye == rh - 1 def _check_near_window(self, pid: str) -> bool: if not self.windows: return True placed = self.placed_items[pid] xb, xe, yb, ye = self._bbox(placed, self.item_defs[pid]) for wx, wy in self.windows: for gy in range(yb, ye + 1): for gx in range(xb, xe + 1): if abs(gx - wx) + abs(gy - wy) <= NEAR_WINDOW_RADIUS: return True return False def _check_near_door(self, pid: str) -> bool: dx, dy = self.door_position placed = self.placed_items[pid] xb, xe, yb, ye = self._bbox(placed, self.item_defs[pid]) for gy in range(yb, ye + 1): for gx in range(xb, xe + 1): if abs(gx - dx) + abs(gy - dy) <= NEAR_DOOR_RADIUS: return True return False def _check_must_face(self, pid: str, target_id: str) -> bool: if target_id not in self.placed_items: return False p = self.placed_items[pid] t = self.placed_items[target_id] pxb, pxe, pyb, pye = self._bbox(p, self.item_defs[pid]) txb, txe, tyb, tye = self._bbox(t, self.item_defs[target_id]) pcx = (pxb + pxe) / 2 pcy = (pyb + pye) / 2 tcx = (txb + txe) / 2 tcy = (tyb + tye) / 2 aligned_x = abs(pcx - tcx) <= 2 aligned_y = abs(pcy - tcy) <= 2 # True if they share one axis (same column or same row) but not both return (aligned_x and not aligned_y) or (aligned_y and not aligned_x) def _check_adjacent_to(self, pid: str, target_id: str) -> bool: if target_id not in self.placed_items: return False t = self.placed_items[target_id] tdef = self.item_defs[target_id] txb, txe, tyb, tye = self._bbox(t, tdef) for nx, ny in self._all_neighbors(pid): if txb <= nx <= txe and tyb <= ny <= tye: return True return False def _check_all_sides_clear(self, pid: str, grid: List[List[bool]]) -> bool: rw, rh = self.room_size placed = self.placed_items[pid] xb, xe, yb, ye = self._bbox(placed, self.item_defs[pid]) top = yb > 0 and any(not grid[yb-1][gx] for gx in range(xb, xe+1)) bottom = ye+1 < rh and any(not grid[ye+1][gx] for gx in range(xb, xe+1)) left = xb > 0 and any(not grid[gy][xb-1] for gy in range(yb, ye+1)) right = xe+1 < rw and any(not grid[gy][xe+1] for gy in range(yb, ye+1)) return top and bottom and left and right def _check_needs_sunlight(self, pid: str, solar_grid: List[List[bool]]) -> bool: placed = self.placed_items[pid] xb, xe, yb, ye = self._bbox(placed, self.item_defs[pid]) for gy in range(yb, ye + 1): for gx in range(xb, xe + 1): if solar_grid[gy][gx]: return True return False def _check_thermal_comfort(self, pid: str, min_v: float, max_v: float, thermal_grid: List[List[float]]) -> bool: placed = self.placed_items[pid] xb, xe, yb, ye = self._bbox(placed, self.item_defs[pid]) total = 0.0 count = 0 for gy in range(yb, ye + 1): for gx in range(xb, xe + 1): total += thermal_grid[gy][gx] count += 1 avg = total / count if count > 0 else 0 return min_v <= avg <= max_v # ─── Master Evaluator ───────────────────────────────────────────────────── def _evaluate(self): self.logs.clear() violations = [] is_incomplete = len(self.unplaced_items) > 0 if is_incomplete: self.success = False self.logs.append(f"Incomplete: {len(self.unplaced_items)} item(s) unplaced: {list(self.unplaced_items.keys())}") grid = self._build_grid() self.solar_grid = self._build_solar_grid() self.thermal_grid = self._build_thermal_grid() visited = self._bfs_reachable(grid) rw, rh = self.room_size # 1. Human Flow for iid in self.placed_items: neighbour_reachable = any( 0 <= nx < rw and 0 <= ny < rh and visited[ny][nx] and not grid[ny][nx] for (nx, ny) in self._all_neighbors(iid) ) if not neighbour_reachable: violations.append(f"FLOW ERROR: Human path to '{iid}' is completely blocked.") # 2. Per-item architectural rules for iid, placed in self.placed_items.items(): item_def = self.item_defs[iid] for rule in item_def.rules: err = None if rule == "must_touch_wall": if not self._check_must_touch_wall(iid): err = f"DESIGN ERROR [{iid.upper()}]: Must be flush against a wall." elif rule == "near_window": if not self._check_near_window(iid): err = f"DESIGN ERROR [{iid.upper()}]: Must be within {NEAR_WINDOW_RADIUS} tiles of a window." elif rule == "near_door": if not self._check_near_door(iid): err = f"DESIGN ERROR [{iid.upper()}]: Must be within {NEAR_DOOR_RADIUS} tiles of the door." elif rule == "all_sides_clear": if not self._check_all_sides_clear(iid, grid): err = f"DESIGN ERROR [{iid.upper()}]: Must have free space on all 4 sides." elif rule == "needs_sunlight": if not self._check_needs_sunlight(iid, self.solar_grid): err = f"DESIGN ERROR [{iid.upper()}]: Blocked from natural sunlight by a tall object." elif rule.startswith("thermal_comfort:"): parts = rule.split(":") if not self._check_thermal_comfort(iid, float(parts[1]), float(parts[2]), self.thermal_grid): err = f"DESIGN ERROR [{iid.upper()}]: Outside thermal comfort limits ({parts[1]} to {parts[2]})." elif rule.startswith("must_face:"): target_id = rule.split(":")[1] if not self._check_must_face(iid, target_id): err = f"DESIGN ERROR [{iid.upper()}]: Must align axis with {target_id.upper()}." elif rule.startswith("adjacent_to:"): target_id = rule.split(":")[1] if not self._check_adjacent_to(iid, target_id): err = f"DESIGN ERROR [{iid.upper()}]: Must be placed adjacent to {target_id.upper()}." if err: violations.append(err) if violations: self.logs.extend(violations) self.success = False elif not is_incomplete: self.logs = ["SUCCESS: All items placed. Human flow verified. All architectural constraints satisfied!"] self.success = True return len(violations) # ─── Step ───────────────────────────────────────────────────────────────── def step(self, action: Action) -> Tuple[Observation, float, bool, dict]: self.step_count += 1 # Dense reward variables reward = 0.0 # Calculate pre-action violations (to reward fixing errors) pre_violations = self._evaluate() if self.placed_items else 0 if action.action_type == "place_item" and action.item: if action.item in self.unplaced_items: item_def = self.unplaced_items[action.item] elif action.item in self.placed_items: item_def = self.item_defs[action.item] del self.placed_items[action.item] else: self.logs = [f"ERROR: '{action.item}' not found in inventory."] return self._get_obs(), -0.5, False, {} xb, xe, yb, ye = self._bbox_action(action, item_def) rw, rh = self.room_size if xb < 0 or xe >= rw or yb < 0 or ye >= rh: self.logs = [f"BOUNDS: '{action.item}' at ({action.x},{action.y}) is outside the room ({rw}x{rh})."] self.unplaced_items[action.item] = item_def reward = -1.0 else: overlap_id = None for oid, op in self.placed_items.items(): oxb, oxe, oyb, oye = self._bbox(op, self.item_defs[oid]) if not (xe < oxb or xb > oxe or ye < oyb or yb > oye): overlap_id = oid break if overlap_id: self.logs = [f"COLLISION: '{action.item}' overlaps with '{overlap_id}'."] self.unplaced_items[action.item] = item_def reward = -1.0 else: if action.item in self.unplaced_items: del self.unplaced_items[action.item] self.placed_items[action.item] = PlacedFurniture( id=action.item, x=action.x, y=action.y, rotation=action.rotation ) # Dense partial progress reward logic post_violations = self._evaluate() if post_violations < pre_violations: reward = 1.0 # Fixed a violation! elif post_violations == 0 and len(self.unplaced_items) < len(self.item_defs): reward = 0.5 # Placed an item cleanly, no violations! else: reward = -0.2 * post_violations # Penalize the current active violations elif action.action_type == "remove_item" and action.item in self.placed_items: self.unplaced_items[action.item] = self.item_defs[action.item] del self.placed_items[action.item] self._evaluate() reward = -0.1 # Slight penalty for backtracking done = self.success or self.step_count >= self.max_steps if self.success: reward = 10.0 # Big sparse bonus for completing the level return self._get_obs(), reward, done, {} def _get_obs(self) -> Observation: return Observation( room_size=self.room_size, door_position=self.door_position, windows=self.windows, unplaced_items=list(self.unplaced_items.values()), placed_items=copy.deepcopy(self.placed_items), solar_grid=copy.deepcopy(self.solar_grid), thermal_grid=copy.deepcopy(self.thermal_grid), feedback_logs=copy.deepcopy(self.logs), step_count=self.step_count, max_steps=self.max_steps, ) def get_episode_result(self) -> EpisodeResult: # Meaningful score mapping (0.0 to 1.0) # 0.0 - 0.5: Partial progress based on placed valid items. # 0.5 - 1.0: Success efficiency. total_items = len(self.item_defs) placed_ratio = len(self.placed_items) / total_items if total_items > 0 else 0 # Penalize for active violations active_violations = len([l for l in self.logs if "ERROR" in l]) valid_placement_score = max(0.0, placed_ratio * 0.5 - (active_violations * 0.05)) if self.success: # The baseline agent places everything perfectly in exactly `total_items` steps. wasted_steps = max(0, self.step_count - total_items) efficiency_bonus = max(0.0, 0.5 - (wasted_steps * 0.05)) score = 0.5 + efficiency_bonus else: score = valid_placement_score # OpenEnv Validator Requirement: Strictly between (0.0, 1.0) final_score = min(0.99, max(0.01, score)) return EpisodeResult( task_id=self.task_id, success=self.success, score=final_score, steps_taken=self.step_count, )