import gymnasium as gym import numpy as np import pygame import math class PotionEnv(gym.Env): metadata = {"render_modes": ["rgb_array"], "render_fps": 30} def __init__(self, max_episode_steps=250, recipe=None): super().__init__() self.max_episode_steps = max_episode_steps # canvas / play area self.width, self.height = 224, 224 # logical grid: 5 cols x 4 rows (playable) self.grid_cols = 5 self.grid_rows = 4 # one extra non-playable row at the top for background / UI self.reserved_top_rows = 1 self.total_rows = self.grid_rows + self.reserved_top_rows self.area_size = 184 self.area_x = (self.width - self.area_size) // 2 self.area_y = (self.height - self.area_size) // 2 self.cell_w = self.area_size / self.grid_cols self.cell_h = self.area_size / self.total_rows # includes reserved row # agent (little witch) self.agent_r = 7 self.speed = 2.5 self.friction = 0.7 self.max_speed = self.speed * 2.0 # cauldron roughly centered (logical coords) self.cauldron_row = 1 # logical row index in [0, grid_rows-1] self.cauldron_col = self.grid_cols // 2 self.cauldron_r = 12 # ingredients (4 total) self.n_ingredients = 4 self.max_recipe_len = 3 # layout (0 = empty, 1 = cauldron, 2-6 = ingredient slots) self.layout_grid = np.zeros( (self.grid_rows, self.grid_cols), dtype=np.int32 ) self.ingredient_id_grid = -np.ones_like(self.layout_grid, dtype=np.int32) self.ingredient_positions = [None] * self.n_ingredients # table bounds for collision & rendering self._compute_table_bounds() self._setup_layout() # runtime state self.ax = self.ay = 0.0 self.vx = self.vy = 0.0 self.carrying = -1 # -1 = nothing, 0..3 = ingredient id # recipe info (in-order) self.base_recipe = list(recipe) if recipe is not None else [0, 2, 3] self.recipe = [] self.collected = [] # bool flags per recipe element self.recipe_step = 0 # index into recipe self.brew_progress = 0.0 self.brew_wrong = 0 self.success = False self.steps = 0 # reward scales self.pick_reward_scale = 0.3 # stage 1: move to correct ingredient self.carry_reward_scale = 0.4 # stage 2: move to cauldron with correct ingredient self.progress_per_step = 0.5 # stage 3: per-step bonus per completed ingredient self.completion_bonus = 10.0 # one-time bonus when recipe is fully completed # sparkles around cauldron when correct ingredient is added self._sparkles = [] # list of dicts: {'start_step', 'ing_id'} self._sparkle_duration = 24 # completion visuals self._completion_step = None # step index when recipe first completed self._completion_fade_frames = 30 self._completion_max_alpha = 70 self._completion_overlay_color = (120, 220, 160) # pygame visuals self._surface = None # final (possibly shaken/flashed) surface self._world_surface = None # base world (no shake/flash) self._font_small = None # shake / flash for wrong ingredient self._shake_frames = 0 self._shake_intensity = 2 self._flash_frames = 0 self._flash_max_frames = 6 # palette self._bg_wall = (245, 236, 222) self._floor_color = (220, 190, 160) self._floor_line = (205, 175, 145) self._table_color = (176, 140, 110) self._table_edge = (140, 105, 80) self._cauldron_body = (40, 40, 55) self._cauldron_rim = (25, 25, 40) self._cauldron_shadow = (30, 25, 35) self._steam_color = (235, 235, 245) self._witch_body = (70, 60, 90) self._witch_accent = (200, 180, 220) self._witch_hat = (55, 40, 80) self._witch_skin = (240, 210, 190) self._witch_boot = (50, 40, 60) # high-contrast ingredient colors (simple bottles) self._ingredient_colors = [ (150, 210, 150), # 0 herb - green (150, 190, 240), # 1 blue vial (230, 170, 230), # 2 crystal - purple/pink (235, 205, 140), # 3 golden spice ] # observation: agent(2) + vel(2) # + carrying (1 + n_ingredients one-hot) # + recipe_ids(3) + recipe_mask(3) + recipe_progress(1) # + layout_grid(5*4) obs_dim = ( 2 + 2 + (1 + self.n_ingredients) + self.max_recipe_len + self.max_recipe_len + 1 + self.grid_rows * self.grid_cols ) self.observation_space = gym.spaces.Box( low=-np.inf, high=np.inf, shape=(obs_dim,), dtype=np.float32 ) # continuous 2D thrust in [-1,1]^2 self.action_space = gym.spaces.Box( low=-np.ones(2, dtype=np.float32), high=np.ones(2, dtype=np.float32), dtype=np.float32, ) # ---------- helpers ---------- def _compute_table_bounds(self): x0 = self.area_x + 4 x1 = self.area_x + self.area_size - 4 # top table (between physical row 1 and 2) top_y0 = self.area_y + self.cell_h + 2 top_y1 = self.area_y + 2 * self.cell_h - 4 # bottom table hugging the bottom row (physical last row) bottom_y0 = self.area_y + self.area_size - self.cell_h + 4 bottom_y1 = self.area_y + self.area_size - 2 self._top_table_bounds = (float(x0), float(top_y0), float(x1), float(top_y1)) self._bottom_table_bounds = (float(x0), float(bottom_y0), float(x1), float(bottom_y1)) def _setup_layout(self): self.layout_grid.fill(0) self.ingredient_id_grid.fill(-1) # cauldron in center (logical coords) self.layout_grid[self.cauldron_row, self.cauldron_col] = 1 self.cauldron_cx, self.cauldron_cy = self._cell_center( self.cauldron_row, self.cauldron_col ) # ingredients: positions = [ (0, 0), # id 0 (0, 4), # id 1 (3, 0), # id 2 (3, 4), # id 3 ] for idx, (r, c) in enumerate(positions): self.layout_grid[r, c] = 2 + idx self.ingredient_id_grid[r, c] = idx self.ingredient_positions[idx] = (r, c) def _clamp_to_bounds(self, x, y, r): min_x = self.area_x + r + 1 max_x = self.area_x + self.area_size - r - 1 min_y = self.area_y + self.reserved_top_rows * self.cell_h + r + 1 max_y = self.area_y + self.area_size - r - 1 x = np.clip(x, min_x, max_x) y = np.clip(y, min_y, max_y) return x, y def _cell_from_pos(self, x, y): if ( x < self.area_x or x >= self.area_x + self.area_size or y < self.area_y or y >= self.area_y + self.area_size ): return None col = int((x - self.area_x) / self.cell_w) row_phys = int((y - self.area_y) / self.cell_h) if ( row_phys < self.reserved_top_rows or row_phys >= self.total_rows ): return None row = row_phys - self.reserved_top_rows # logical row col = max(0, min(self.grid_cols - 1, col)) row = max(0, min(self.grid_rows - 1, row)) return row, col def _cell_center(self, row, col): cx = self.area_x + (col + 0.5) * self.cell_w cy = self.area_y + (self.reserved_top_rows + row + 0.5) * self.cell_h return cx, cy def _push_out_of_rect(self, x0, y0, x1, y1): ax, ay, r = self.ax, self.ay, self.agent_r nx = min(max(ax, x0), x1) ny = min(max(ay, y0), y1) dx = ax - nx dy = ay - ny dist = math.hypot(dx, dy) if dist < r: if dist < 1e-6: mid_y = 0.5 * (y0 + y1) if ay < mid_y: self.ay = y0 - r - 1 else: self.ay = y1 + r + 1 else: overlap = r - dist + 0.5 self.ax += (dx / dist) * overlap self.ay += (dy / dist) * overlap self.vx *= 0.2 self.vy *= 0.2 def _avoid_cauldron(self): dx = self.ax - self.cauldron_cx dy = self.ay - self.cauldron_cy dist = math.hypot(dx, dy) min_dist = self.cauldron_r + self.agent_r + 2.0 if dist < 1e-6: dx, dy = 1.0, 0.0 dist = 1.0 if dist < min_dist: scale = min_dist / dist self.ax = self.cauldron_cx + dx * scale self.ay = self.cauldron_cy + dy * scale self.vx *= 0.0 self.vy *= 0.0 def _avoid_tables(self): self._push_out_of_rect(*self._top_table_bounds) self._push_out_of_rect(*self._bottom_table_bounds) def _emit_sparkles(self, ing_id): self._sparkles.append({"start_step": self.steps, "ing_id": int(ing_id)}) def _trigger_wrong_fx(self): self._shake_frames = max(self._shake_frames, 10) self._flash_frames = self._flash_max_frames def _add_to_cauldron(self, ing_id): correct = False wrong = False if self.recipe_step >= len(self.recipe): wrong = True self.brew_wrong += 1 return correct, wrong required = self.recipe[self.recipe_step] if ing_id == required: self.collected[self.recipe_step] = True self.recipe_step += 1 self.brew_progress = self.recipe_step / float(len(self.recipe)) correct = True self._emit_sparkles(ing_id) # mark recipe completion moment if ( self.recipe_step == len(self.recipe) and self._completion_step is None ): self._completion_step = self.steps elif ing_id in self.recipe[: self.recipe_step]: correct = False wrong = False else: wrong = True self.brew_wrong += 1 return correct, wrong def _get_obs(self): ax_n = self.ax / self.width ay_n = self.ay / self.height vx_n = np.clip(self.vx / self.max_speed, -1.0, 1.0) vy_n = np.clip(self.vy / self.max_speed, -1.0, 1.0) carry_oh = np.zeros(1 + self.n_ingredients, dtype=np.float32) if self.carrying < 0: carry_oh[0] = 1.0 else: carry_oh[1 + int(self.carrying)] = 1.0 sentinel = float(self.n_ingredients) # "none" id recipe_ids = np.full(self.max_recipe_len, sentinel, dtype=np.float32) recipe_mask = np.zeros(self.max_recipe_len, dtype=np.float32) for i in range(min(len(self.recipe), self.max_recipe_len)): recipe_ids[i] = float(self.recipe[i]) recipe_mask[i] = 1.0 denom = max(1.0, float(self.n_ingredients - 1)) recipe_ids_norm = recipe_ids / denom if len(self.recipe) > 0: step_norm = self.recipe_step / float(len(self.recipe)) else: step_norm = 0.0 layout_flat = (self.layout_grid.flatten().astype(np.float32) / 7.0).tolist() obs = np.array( [ax_n, ay_n, vx_n, vy_n] + carry_oh.tolist() + recipe_ids_norm.tolist() + recipe_mask.tolist() + [step_norm] + layout_flat, dtype=np.float32, ) return obs # ---------- gym API ---------- def reset(self, *, seed=None, options=None): super().reset(seed=seed) # random starting cell among interior rows (1 and 2) that are empty candidates = [] for r in range(self.grid_rows): for c in range(self.grid_cols): if self.layout_grid[r, c] == 0 and r in (1, 2): candidates.append((r, c)) assert candidates, "No empty cells for spawn." idx = self.np_random.integers(0, len(candidates)) start_row, start_col = candidates[idx] self.ax, self.ay = self._cell_center(start_row, start_col) self.ax, self.ay = self._clamp_to_bounds(self.ax, self.ay, self.agent_r) self.vx = self.vy = 0.0 self.carrying = -1 self.recipe = list(self.base_recipe) self.collected = [False] * len(self.recipe) self.recipe_step = 0 self.brew_progress = 0.0 self.brew_wrong = 0 self.success = False self._sparkles.clear() self._shake_frames = 0 self._flash_frames = 0 # reset completion visuals self._completion_step = None self.steps = 0 return self._get_obs(), {} def step(self, action): action = np.clip(np.asarray(action, dtype=np.float32), -1.0, 1.0) # velocity with light inertia self.vx = self.vx * self.friction + float(action[0]) * self.speed self.vy = self.vy * self.friction + float(action[1]) * self.speed # clamp max speed v_norm = math.hypot(self.vx, self.vy) if v_norm > self.max_speed: scale = self.max_speed / (v_norm + 1e-8) self.vx *= scale self.vy *= scale # move & clamp nx = self.ax + self.vx ny = self.ay + self.vy self.ax, self.ay = self._clamp_to_bounds(nx, ny, self.agent_r) # avoid obstacles self._avoid_cauldron() self._avoid_tables() self.ax, self.ay = self._clamp_to_bounds(self.ax, self.ay, self.agent_r) ingredient_picked = False correct_add = False wrong_add = False completed_now = False # remaining required ingredients (for info, not reward directly) remaining_ids = [ ing for ing, done in zip(self.recipe, self.collected) if not done ] # distances for logging / interactions dist_c = math.hypot(self.ax - self.cauldron_cx, self.ay - self.cauldron_cy) # front-of-cauldron delivery point deliver_tx = self.cauldron_cx deliver_ty = self.cauldron_cy + self.cauldron_r * 0.7 dist_deliver = math.hypot(self.ax - deliver_tx, self.ay - deliver_ty) deliver_radius = self.cauldron_r + self.agent_r + 7.0 # --- ingredient pickup (no direct reward here, just state) --- if self.carrying == -1: best_ing = -1 best_dist = float("inf") pickup_radius = max(self.cell_w, self.cell_h) * 0.8 for ing_id, pos in enumerate(self.ingredient_positions): r_i, c_i = pos cx, cy = self._cell_center(r_i, c_i) d = math.hypot(self.ax - cx, self.ay - cy) if d < pickup_radius and d < best_dist: best_ing = ing_id best_dist = d if best_ing >= 0: self.carrying = int(best_ing) ingredient_picked = True # --- dropping into cauldron --- if dist_deliver <= deliver_radius: if self.carrying != -1: ing_id = self.carrying self.carrying = -1 prev_step = self.recipe_step # for completion detection correct, wrong = self._add_to_cauldron(ing_id) correct_add = correct wrong_add = wrong if wrong_add: # purely visual feedback, no reward effect self._trigger_wrong_fx() # did we just finish the entire recipe this step? if (correct_add and prev_step < len(self.recipe) and self.recipe_step == len(self.recipe)): completed_now = True self.success = True r_pick = 0.0 r_carry = 0.0 # next required ingredient id (None if recipe complete) if self.recipe_step < len(self.recipe): next_ing_id = self.recipe[self.recipe_step] else: next_ing_id = None dist_to_next = None dist_norm_target = 1.0 # for logging: normalized distance to current target if next_ing_id is not None: # world position of next required ingredient r_next, c_next = self.ingredient_positions[next_ing_id] next_x, next_y = self._cell_center(r_next, c_next) dist_to_next = math.hypot(self.ax - next_x, self.ay - next_y) dist_next_norm = np.clip(dist_to_next / self.area_size, 0.0, 1.0) # stage 1: moving toward next ingredient (only if not carrying anything) if self.carrying == -1: r_pick = self.pick_reward_scale * (1.0 - dist_next_norm) dist_norm_target = dist_next_norm # stage 2: carrying correct ingredient to cauldron elif self.carrying == next_ing_id: dist_deliver_norm = np.clip(dist_deliver / self.area_size, 0.0, 1.0) # constant = max of stage 1 so there is no drop at pickup r_carry = self.pick_reward_scale + self.carry_reward_scale * (1.0 - dist_deliver_norm) dist_norm_target = dist_deliver_norm else: # carrying wrong ingredient -> no stage 1/2 shaping dist_norm_target = 1.0 else: # recipe complete: no stage 1/2 shaping, only progress bonus below dist_norm_target = 1.0 # stage 3: per-step progress bonus (k = recipe_step) r_progress = self.recipe_step * self.progress_per_step reward = r_pick + r_carry + r_progress # bonus for completing the recipe if self.success: reward = self.completion_bonus self.steps += 1 terminated = False truncated = self.steps >= self.max_episode_steps info = { "pickup": bool(ingredient_picked), "correct_add": bool(correct_add), "wrong_add": bool(wrong_add), "recipe": list(self.recipe), "recipe_step": int(self.recipe_step), "carrying": int(self.carrying), "brew_progress": float(self.brew_progress), "distance_to_target": float(dist_norm_target), "dist_to_next_ingredient": float(dist_to_next if dist_to_next is not None else -1.0), "dist_to_cauldron_center": float(dist_c), "dist_to_delivery_point": float(dist_deliver), "r_pick": float(r_pick), "r_carry": float(r_carry), "r_progress": float(r_progress), "completed_now": bool(completed_now), "success": bool(self.success), } return self._get_obs(), float(reward), terminated, truncated, info # ---------- rendering ---------- def _liquid_color(self): t = self.steps * 0.12 p = self.brew_progress r = 80 + int(60 * p + 20 * math.sin(t)) g = 140 + int(40 * (1.0 - p) + 10 * math.cos(t * 0.7)) b = 170 + int(30 * (p**2) + 15 * math.sin(t * 1.3)) r = max(0, min(255, r)) g = max(0, min(255, g)) b = max(0, min(255, b)) return (r, g, b) def _draw_bottle(self, surf, ing_id, cx, cy): base_color = self._ingredient_colors[ing_id] bw = int(self.cell_w * 0.26) bh = int(self.cell_h * 0.6) x = int(cx - bw / 2) y = int(cy - bh / 2) + 2 bottle_rect = pygame.Rect(x, y + 4, bw, int(bh * 0.7)) neck_w = int(bw * 0.6) neck_h = int(bh * 0.25) neck_rect = pygame.Rect( int(cx - neck_w / 2), y, neck_w, neck_h ) pygame.draw.rect(surf, base_color, bottle_rect, border_radius=4) pygame.draw.rect(surf, base_color, neck_rect, border_radius=3) highlight = bottle_rect.inflate(-bw // 2, -bottle_rect.height // 2) pygame.draw.rect(surf, (255, 255, 255), highlight, 1, border_radius=3) def _draw_cauldron(self, surf): cx, cy = int(self.cauldron_cx), int(self.cauldron_cy) shadow_rect = pygame.Rect( int(cx - self.cauldron_r * 1.4), int(cy + self.cauldron_r * 0.7), int(self.cauldron_r * 2.8), int(self.cauldron_r * 0.7), ) pygame.draw.ellipse(surf, self._cauldron_shadow, shadow_rect) body_rect = pygame.Rect( int(cx - self.cauldron_r), int(cy - self.cauldron_r * 0.1), int(self.cauldron_r * 2), int(self.cauldron_r * 1.9), ) pygame.draw.ellipse(surf, self._cauldron_body, body_rect) rim_rect = pygame.Rect( int(cx - self.cauldron_r * 1.1), int(cy - self.cauldron_r * 0.7), int(self.cauldron_r * 2.2), int(self.cauldron_r * 0.9), ) pygame.draw.ellipse(surf, self._cauldron_rim, rim_rect) pygame.draw.ellipse(surf, (10, 10, 20), rim_rect, 1) liquid_rect = rim_rect.inflate( -int(self.cauldron_r * 0.5), -int(self.cauldron_r * 0.5) ) pygame.draw.ellipse(surf, self._liquid_color(), liquid_rect) lx, ly = liquid_rect.center inner_r = liquid_rect.width * 0.22 for i in range(3): phase = self.steps * 0.25 + i * 2.0 jitter = 1.5 * math.sin(self.steps * 0.4 + i) bx = int(lx + (inner_r * 0.5 + jitter) * math.cos(phase)) by = int(ly + (inner_r * 0.3 + jitter) * math.sin(phase)) pygame.draw.circle(surf, (245, 250, 255), (bx, by), 2) for i in range(3): offset = (i - 1) * 6 sway = int(2 * math.sin(self.steps * 0.2 + i)) rise = int(2 * math.cos(self.steps * 0.3 + i)) steam_rect = pygame.Rect( cx - 3 + offset + sway, int(liquid_rect.top - 16 + rise), 6, 18, ) pygame.draw.ellipse(surf, self._steam_color, steam_rect) def _draw_sparkles(self, surf): if not self._sparkles: return cx, cy = int(self.cauldron_cx), int(self.cauldron_cy) updated = [] for s in self._sparkles: age = self.steps - s["start_step"] if age < 0 or age > self._sparkle_duration: continue t = age / float(self._sparkle_duration) ing_id = s["ing_id"] br, bg, bb = self._ingredient_colors[ing_id] br = min(255, int(br + 70)) bg = min(255, int(bg + 70)) bb = min(255, int(bb + 70)) alpha_scale = max(0.0, 1.0 - t) color = ( int(br * alpha_scale + 255 * (1.0 - alpha_scale) * 0.1), int(bg * alpha_scale + 255 * (1.0 - alpha_scale) * 0.1), int(bb * alpha_scale + 255 * (1.0 - alpha_scale) * 0.1), ) radius = self.cauldron_r + 6 + t * 6.0 dot_r = 3 for i in range(6): angle = 2.0 * math.pi * i / 6.0 x = int(cx + radius * math.cos(angle)) y = int(cy + radius * math.sin(angle) * 0.8) pygame.draw.circle(surf, color, (x, y), dot_r) updated.append(s) self._sparkles = updated def _draw_agent(self, surf): ax = int(self.ax) ay = int(self.ay) r = self.agent_r body_rect = pygame.Rect(ax - r, ay - r + 3, 2 * r, 2 * r + 2) pygame.draw.ellipse(surf, self._witch_body, body_rect) accent_rect = body_rect.inflate(-int(0.6 * r), -int(0.4 * r)) pygame.draw.ellipse(surf, self._witch_accent, accent_rect) speed = math.hypot(self.vx, self.vy) phase = self.steps * 0.35 amp = min(3.0, (speed / (self.max_speed + 1e-6)) * 3.0) foot_y_base = ay + r + 2 offset = int(amp * math.sin(phase)) foot_w = int(r * 0.9) foot_h = int(r * 0.4) left_foot = pygame.Rect(ax - foot_w, foot_y_base + offset, foot_w, foot_h) right_foot = pygame.Rect(ax, foot_y_base - offset, foot_w, foot_h) pygame.draw.ellipse(surf, self._witch_boot, left_foot) pygame.draw.ellipse(surf, self._witch_boot, right_foot) head_r = int(r * 0.7) head_center = (ax, ay - r) pygame.draw.circle(surf, self._witch_skin, head_center, head_r) brim_rect = pygame.Rect(ax - r, head_center[1] + 1, 2 * r, 3) pygame.draw.rect(surf, self._witch_hat, brim_rect) hat_pts = [ (ax, head_center[1] - head_r - 2), (ax - head_r, head_center[1] + 1), (ax + head_r, head_center[1] + 1), ] pygame.draw.polygon(surf, self._witch_hat, hat_pts) eye_dx = int(head_r * 0.4) eye_y = head_center[1] - int(head_r * 0.1) pygame.draw.circle(surf, (20, 20, 30), (head_center[0] - eye_dx, eye_y), 1) pygame.draw.circle(surf, (20, 20, 30), (head_center[0] + eye_dx, eye_y), 1) if self.carrying >= 0: orb_color = self._ingredient_colors[self.carrying] orb_center = (ax, head_center[1] - head_r - 6) pygame.draw.circle(surf, orb_color, orb_center, 4) pygame.draw.circle(surf, (255, 255, 255), orb_center, 4, 1) def _draw_recipe_ui(self, surf): if self._font_small is None: return panel_rect = pygame.Rect(8, 8, 82, 40) # panel background pygame.draw.rect(surf, (250, 245, 240), panel_rect, border_radius=6) # border color depends on completion if len(self.recipe) > 0 and self.recipe_step == len(self.recipe): border_color = (90, 180, 110) border_width = 2 else: border_color = (210, 180, 150) border_width = 1 pygame.draw.rect(surf, border_color, panel_rect, border_width, border_radius=6) label = self._font_small.render("Recipe", True, (80, 60, 50)) surf.blit(label, (panel_rect.x + 6, panel_rect.y + 4)) cx0 = panel_rect.x + 18 cy = panel_rect.y + 24 step_spacing = 22 for i in range(self.max_recipe_len): cx = cx0 + i * step_spacing if i < len(self.recipe): ing_id = self.recipe[i] base_color = self._ingredient_colors[ing_id] pygame.draw.circle(surf, base_color, (cx, cy), 6) if i < self.recipe_step: border = (80, 150, 90) else: border = (120, 100, 90) pygame.draw.circle(surf, border, (cx, cy), 6, 1) else: pygame.draw.circle(surf, (210, 210, 210), (cx, cy), 5, 1) def render(self): if self._surface is None: pygame.init() self._surface = pygame.Surface((self.width, self.height)) if self._world_surface is None: self._world_surface = pygame.Surface((self.width, self.height)) if self._font_small is None: pygame.font.init() self._font_small = pygame.font.SysFont("arial", 12) world = self._world_surface world.fill(self._bg_wall) floor_y0 = self.area_y + self.reserved_top_rows * self.cell_h floor_h = self.grid_rows * self.cell_h floor_rect = pygame.Rect( self.area_x, int(floor_y0), self.area_size, int(floor_h), ) pygame.draw.rect(world, self._floor_color, floor_rect) for c in range(self.grid_cols + 1): x = int(self.area_x + c * self.cell_w) pygame.draw.line( world, self._floor_line, (x, floor_y0), (x, floor_y0 + floor_h), 1, ) for r in range(self.grid_rows + 1): y = int(floor_y0 + r * self.cell_h) pygame.draw.line( world, self._floor_line, (self.area_x, y), (self.area_x + self.area_size, y), 1, ) x0, y0, x1, y1 = self._top_table_bounds top_rect = pygame.Rect(int(x0), int(y0), int(x1 - x0), int(y1 - y0)) x0b, y0b, x1b, y1b = self._bottom_table_bounds bottom_rect = pygame.Rect(int(x0b), int(y0b), int(x1b - x0b), int(y1b - y0b)) pygame.draw.rect(world, self._table_color, top_rect, border_radius=4) pygame.draw.rect(world, self._table_color, bottom_rect, border_radius=4) pygame.draw.rect(world, self._table_edge, top_rect, 1, border_radius=4) pygame.draw.rect(world, self._table_edge, bottom_rect, 1, border_radius=4) for ing_id in range(self.n_ingredients): row, col = self.ingredient_positions[ing_id] cx, cy = self._cell_center(row, col) self._draw_bottle(world, ing_id, int(cx), int(cy)) self._draw_cauldron(world) self._draw_sparkles(world) self._draw_agent(world) self._draw_recipe_ui(world) surf = self._surface dx = dy = 0 if self._shake_frames > 0: dx = self.np_random.integers(-self._shake_intensity, self._shake_intensity + 1) dy = self.np_random.integers(-self._shake_intensity, self._shake_intensity + 1) self._shake_frames -= 1 surf.fill(self._bg_wall) surf.blit(world, (dx, dy)) # red flash for wrong ingredient if self._flash_frames > 0: alpha = int(120 * (self._flash_frames / float(self._flash_max_frames))) overlay = pygame.Surface((self.width, self.height)) overlay.fill((255, 60, 60)) overlay.set_alpha(alpha) surf.blit(overlay, (0, 0)) self._flash_frames -= 1 # green success overlay: fades in once recipe is complete, then stays if self._completion_step is not None: age = max(0, self.steps - self._completion_step) t = min(1.0, age / float(self._completion_fade_frames)) alpha = int(self._completion_max_alpha * t) if alpha > 0: overlay = pygame.Surface((self.width, self.height)) overlay.fill(self._completion_overlay_color) overlay.set_alpha(alpha) surf.blit(overlay, (0, 0)) arr = np.transpose(np.array(pygame.surfarray.pixels3d(surf)), (1, 0, 2)).copy() return arr def close(self): if self._surface is not None: pygame.quit() self._surface = None self._world_surface = None # --- recipe-specific variations --- class PotionRecipe1Env(PotionEnv): def __init__(self, max_episode_steps=250): super().__init__(max_episode_steps=max_episode_steps, recipe=[0, 2, 3]) class PotionRecipe2Env(PotionEnv): def __init__(self, max_episode_steps=250): super().__init__(max_episode_steps=max_episode_steps, recipe=[1, 3, 0]) class PotionRecipe3Env(PotionEnv): def __init__(self, max_episode_steps=250): super().__init__(max_episode_steps=max_episode_steps, recipe=[2, 3, 1]) class PotionRecipe4Env(PotionEnv): def __init__(self, max_episode_steps=250): super().__init__(max_episode_steps=max_episode_steps, recipe=[3, 1, 2]) class PotionRecipe5Env(PotionEnv): def __init__(self, max_episode_steps=250): super().__init__(max_episode_steps=max_episode_steps, recipe=[3, 2, 1]) class PotionRecipe6Env(PotionEnv): def __init__(self, max_episode_steps=250): super().__init__(max_episode_steps=max_episode_steps, recipe=[2, 1, 3]) class PotionRecipe7Env(PotionEnv): def __init__(self, max_episode_steps=250): super().__init__(max_episode_steps=max_episode_steps, recipe=[0, 1, 2]) class PotionRecipe8Env(PotionEnv): def __init__(self, max_episode_steps=250): super().__init__(max_episode_steps=max_episode_steps, recipe=[0, 1]) class PotionRecipe9Env(PotionEnv): def __init__(self, max_episode_steps=250): super().__init__(max_episode_steps=max_episode_steps, recipe=[3, 2]) class PotionRecipe10Env(PotionEnv): def __init__(self, max_episode_steps=250): super().__init__(max_episode_steps=max_episode_steps, recipe=[2, 0])