ashucode/metaxhuggingfacehackathon

1import random2from typing import Optional, Tuple, Dict, Any, List3 4import numpy as np5 6try:7    import gymnasium as gym8    from gymnasium import spaces9    GYM_AVAILABLE = True10except ImportError:11    GYM_AVAILABLE = False12    class _Discrete:13        def __init__(self, n): self.n = n14        def sample(self): return random.randint(0, self.n - 1)15    class _Box:16        def __init__(self, **kw): pass17    class _spaces:18        Discrete = _Discrete; Box = _Box19    class _gym:20        Env = object; spaces = _spaces()21    gym = _gym(); spaces = _spaces()22 23 24 25 26PRODUCT_CATALOGUE: Dict[str, List[Dict]] = {27    "> Dairy":     [{"name": "Whole Milk",      "aisle": "A1", "reward": 10},28                   {"name": "Cheddar Cheese",   "aisle": "A1", "reward": 15},29                   {"name": "Curd",             "aisle": "A1", "reward": 12}],30    "> Bakery":    [{"name": "Bread",           "aisle": "A2", "reward": 10},31                   {"name": "Cake",             "aisle": "A2", "reward":  8},32                   {"name": "Pastry",           "aisle": "A2", "reward":  7}],33    "> Fruits":    [{"name": "Apples",          "aisle": "A3", "reward": 10},34                   {"name": "Strawberry",       "aisle": "A3", "reward": 12},35                   {"name": "Kiwi",             "aisle": "A3", "reward": 15}],36    "> Frozen":    [{"name": "Frozen Pizza",    "aisle": "A4", "reward":  8},37                   {"name": "Ice Cream Tub",    "aisle": "A4", "reward": 10},38                   {"name": "Veggie Burger",    "aisle": "A4", "reward": 11}],39    "> Beverages": [{"name": "Orange Juice",    "aisle": "A5", "reward":  9},40                   {"name": "Bisleri",          "aisle": "A5", "reward":  6},41                   {"name": "Cold Coffee",      "aisle": "A5", "reward": 14}],42    "> Snacks":    [{"name": "Lays Chips",      "aisle": "A6", "reward":  7},43                   {"name": "Mix Namkeen",      "aisle": "A6", "reward":  9},44                   {"name": "Dark Chocolate",   "aisle": "A6", "reward": 13}],45}46 47ALL_PRODUCTS: Dict[str, Dict] = {}48for _cat, _items in PRODUCT_CATALOGUE.items():49    for _item in _items:50        ALL_PRODUCTS[_item["name"].lower()] = {**_item, "category": _cat}51 52AISLE_GRID_POS: Dict[str, Tuple[int, int]] = {53    "A1": (1, 1), "A2": (3, 1), "A3": (5, 1),54    "A4": (1, 5), "A5": (3, 5), "A6": (5, 5),55}56AISLE_SYMBOL: Dict[str, str] = {57    "A1": "D", "A2": "B", "A3": "P",58    "A4": "F", "A5": "V", "A6": "S",59}60 61LEVEL_CONFIG: Dict[str, Dict] = {62    "easy": {63        "label":            "EASY",64        "description":      "1 product · 200 steps · 2× rewards · walk to checkout",65        "num_products":     1,66        "max_steps":        200,67        "reward_mult":      2.0,68        "step_penalty":     -0.5,69        "wrong_penalty":    0.0,70        "shaping_scale":    1.0,71        "completion_bonus": 20.0,72        "checkout_bonus":   15.0,73        "closest_rule":     False,74    },75    "medium": {76        "label":            "MEDIUM",77        "description":      "3 products · 80 steps · any order · walk to checkout",78        "num_products":     3,79        "max_steps":        80,80        "reward_mult":      1.0,81        "step_penalty":     -1.0,82        "wrong_penalty":    -5.0,83        "shaping_scale":    1.5,84        "completion_bonus": 35.0,85        "checkout_bonus":   20.0,86        "closest_rule":     False,87    },88    "hard": {89        "label":            "HARD",90        "description":      "4 products · 50 steps · CLOSEST first · walk to checkout",91        "num_products":     4,92        "max_steps":        50,93        "reward_mult":      1.5,94        "step_penalty":     -2.0,95        "wrong_penalty":    -15.0,96        "shaping_scale":    2.0,97        "completion_bonus": 50.0,98        "checkout_bonus":   30.0,99        "closest_rule":     True,100    },101}102 103ACTION_DELTAS  = {0: (-1, 0), 1: (1, 0), 2: (0, -1), 3: (0, 1)}104ACTION_NAMES   = ["UP", "DOWN", "LEFT", "RIGHT", "COLLECT"]105COLLECT_ACTION = 4106N_ACTIONS      = 5107MAX_PRODUCTS   = 4108OBS_LEN        = 4 + 3 * MAX_PRODUCTS  # 16109 110 111 112def _max_possible_reward(cfg: Dict, products: List[Dict]) -> float:113    """114    Theoretical upper bound for a perfect episode (no step penalties,115    all collect bonuses + completion_bonus + checkout_bonus).116    Used to normalise the final score to [0, 1].117    """118    item_rewards = sum(p["reward"] * cfg["reward_mult"] for p in products)119    return item_rewards + cfg["completion_bonus"] + cfg["checkout_bonus"]120 121 122 123class SupermarketEnv:124    GRID_ROWS = 8125    GRID_COLS = 7126    SPAWN_POS = (0, 3)127    EXIT_POS  = (7, 3)128 129    def __init__(self, level: str = "easy",130                 products: Optional[List[str]] = None,131                 render_mode: Optional[str] = None):132        if level not in LEVEL_CONFIG:133            raise ValueError(f"level must be one of {list(LEVEL_CONFIG)}")134        self.level           = level135        self.cfg             = LEVEL_CONFIG[level]136        self._fixed_products = products137        self.render_mode     = render_mode138 139        if GYM_AVAILABLE:140            self.observation_space = spaces.Box(141                low=0.0, high=1.0, shape=(OBS_LEN,), dtype=np.float32)142            self.action_space = spaces.Discrete(N_ACTIONS)143        else:144            self.action_space = _Discrete(N_ACTIONS)145 146        self._agent_pos:       Tuple[int, int] = self.SPAWN_POS147        self._target_products: List[Dict]      = []148        self._collected:       List[bool]      = []149        self._phase:           int             = 0150        self._steps_taken:     int             = 0151        self._closest_idx:     int             = 0152        self._prev_dist:       int             = 0153        self._total_reward:    float           = 0.0154        self._reward_log:      List[Dict]      = []155        self._task_status:     str             = "In-Progress"156        self._max_reward:      float           = 1.0   # set properly after reset157 158    159    def reset(self, *, seed=None, options=None):160        if seed is not None:161            random.seed(seed); np.random.seed(seed)162        if options:163            if "level" in options:164                self.level = options["level"]165                self.cfg   = LEVEL_CONFIG[self.level]166            if "products" in options:167                self._fixed_products = options["products"]168 169        self._target_products = self._resolve_products()170        self._agent_pos    = self.SPAWN_POS171        self._collected    = [False] * len(self._target_products)172        self._phase        = 0173        self._steps_taken  = 0174        self._total_reward = 0.0175        self._reward_log   = []176        self._task_status  = "In-Progress"177        self._closest_idx  = self._compute_closest_idx()178        self._prev_dist    = self._dist_to_goal()179        self._max_reward   = _max_possible_reward(self.cfg, self._target_products)180 181        return self._build_obs(), self._build_info("reset", 0.0)182 183    def step(self, action: int):184        reward     = 0.0185        terminated = False186        truncated  = False187        event      = ""188 189        reward            += self.cfg["step_penalty"]190        self._steps_taken += 1191 192        if action in ACTION_DELTAS:193            dr, dc = ACTION_DELTAS[action]194            nr     = self._agent_pos[0] + dr195            nc     = self._agent_pos[1] + dc196 197            if 0 <= nr < self.GRID_ROWS and 0 <= nc < self.GRID_COLS:198                self._agent_pos   = (nr, nc)199                new_dist          = self._dist_to_goal()200                reward           += self.cfg["shaping_scale"] * (self._prev_dist - new_dist)201                self._prev_dist   = new_dist202                self._closest_idx = self._compute_closest_idx()203                event = f"MOVE_{ACTION_NAMES[action]} → ({nr},{nc})"204 205                if self._phase == 1 and self._agent_pos == self.EXIT_POS:206                    reward           += self.cfg["checkout_bonus"]207                    self._task_status = "Success"208                    terminated        = True209                    event            += f" | REACHED COUNTER ✅ +{self.cfg['checkout_bonus']}"210            else:211                reward -= 1.0212                event   = f"WALL_HIT (tried {ACTION_NAMES[action]})"213 214        elif action == COLLECT_ACTION:215            if self._phase == 0:216                r, event  = self._try_collect()217                reward   += r218                self._closest_idx = self._compute_closest_idx()219 220                if all(self._collected):221                    reward         += self.cfg["completion_bonus"]222                    self._phase     = 1223                    self._prev_dist = self._dist_to_goal()224                    event          += (f" | ALL COLLECTED ✅ +{self.cfg['completion_bonus']}"225                                       f" → RETURN TO COUNTER at {self.EXIT_POS}")226 227                    if self._agent_pos == self.EXIT_POS:228                        reward           += self.cfg["checkout_bonus"]229                        self._task_status = "Success"230                        terminated        = True231                        event            += f" | INSTANT CHECKOUT ✅ +{self.cfg['checkout_bonus']}"232                else:233                    self._prev_dist = self._dist_to_goal()234            else:235                event = "COLLECT ignored — already in exit phase, head to counter"236 237        if not terminated and self._steps_taken >= self.cfg["max_steps"]:238            reward           -= 20.0239            self._task_status = "Failed"240            truncated         = True241            event            += " | BUDGET_EXCEEDED −20"242 243        self._total_reward += reward244        self._append_log(action, reward, event)245        info = self._build_info(event, reward)246        return self._build_obs(), float(reward), terminated, truncated, info247 248   def normalised_score(self) -> float:249        raw = max(0.0, self._total_reward)250        score = round(min(raw / self._max_reward, 1.0), 6) if self._max_reward > 0 else 0.0251        return max(0.001, min(0.999, score))252    253    def action_masks(self) -> np.ndarray:254        mask     = np.ones(N_ACTIONS, dtype=bool)255        on_aisle = self._agent_pos in AISLE_GRID_POS.values()256        if not on_aisle or self._phase == 1:257            mask[COLLECT_ACTION] = False258        else:259 260            aisle = next(261                (code for code, pos in AISLE_GRID_POS.items() if self._agent_pos == pos),262                None,263            )264 265            pending = [266                i for i, (p, c) in enumerate(zip(self._target_products, self._collected))267                if not c and p["aisle"] == aisle268            ]269            if not pending:270                mask[COLLECT_ACTION] = False271            elif self.cfg["closest_rule"]:272                # On HARD level: only allow COLLECT at the closest aisle273                req_aisle = self._target_products[self._closest_idx]["aisle"]274                if aisle != req_aisle:275                    mask[COLLECT_ACTION] = False276        return mask277 278    def get_state_key(self) -> tuple:279        return (280            self._agent_pos[0],281            self._agent_pos[1],282            self._phase,283            tuple(self._collected),284            self._closest_idx,285        )286 287    def render(self):288        if self.render_mode != "ansi":289            return None290        grid = [["  · "] * self.GRID_COLS for _ in range(self.GRID_ROWS)]291        for code, (r, c) in AISLE_GRID_POS.items():292            sym = AISLE_SYMBOL[code]293            is_pending = any(not self._collected[i]294                             and self._target_products[i]["aisle"] == code295                             for i in range(len(self._target_products)))296            is_closest = (not all(self._collected)297                          and self._target_products[self._closest_idx]["aisle"] == code298                          and self._phase == 0)299            if is_closest:300                grid[r][c] = " [★] "301            elif is_pending:302                grid[r][c] = f" [{sym}] "303            else:304                grid[r][c] = f"  {sym.lower()}  "305 306        er, ec = self.EXIT_POS307        grid[er][ec] = "[CTR]" if self._phase == 1 else " ctr "308 309        ar, ac = self._agent_pos310        grid[ar][ac] = "  @  "311 312        phase_str = ("Phase 0: COLLECTING PRODUCTS" if self._phase == 0313                     else "Phase 1: RETURN TO COUNTER")314        print()315        print(f"  ╔═══ {self.level.upper()} ═══ Step {self._steps_taken}/{self.cfg['max_steps']} "316              f"═══ Reward: {self._total_reward:.1f} ═══ {self._task_status} ╗")317        print(f"  ║  {phase_str:<55}║")318        print(f"  ╠{'═'*62}╣")319        header = "   " + "".join(f"  {c}   " for c in range(self.GRID_COLS))320        print(f"  ║ {header}║")321        print(f"  ╠{'═'*62}╣")322        for r, row in enumerate(grid):323            print(f"  ║ {r} {''.join(row)} ║")324        print(f"  ╚{'═'*62}╝")325 326    def close(self): pass327 328    @property329    def task_status(self):     return self._task_status330    @property331    def reward_log(self):      return self._reward_log332    @property333    def total_reward(self):    return self._total_reward334    @property335    def target_names(self):    return [p["name"] for p in self._target_products]336    @property337    def collected_names(self): return [p["name"] for p, c in338                                       zip(self._target_products, self._collected) if c]339 340 341    def _manhattan(self, cell: Tuple[int, int]) -> int:342        return abs(self._agent_pos[0] - cell[0]) + abs(self._agent_pos[1] - cell[1])343 344    def _dist_to_goal(self) -> int:345        if self._phase == 1 or all(self._collected):346            return self._manhattan(self.EXIT_POS)347        return self._manhattan(AISLE_GRID_POS[self._target_products[self._closest_idx]["aisle"]])348 349    def _compute_closest_idx(self) -> int:350        best_i, best_d = 0, float("inf")351        for i, (p, c) in enumerate(zip(self._target_products, self._collected)):352            if c:353                continue354            d = self._manhattan(AISLE_GRID_POS[p["aisle"]])355            if d < best_d:356                best_d, best_i = d, i357        return best_i358 359    def _try_collect(self) -> Tuple[float, str]:360        aisle = next(361            (code for code, pos in AISLE_GRID_POS.items() if self._agent_pos == pos),362            None,363        )364        if aisle is None:365            return self.cfg["wrong_penalty"], "COLLECT_FAIL: not standing on an aisle"366 367        candidates = [368            i for i, (p, c) in enumerate(zip(self._target_products, self._collected))369            if not c and p["aisle"] == aisle370        ]371        if not candidates:372            return self.cfg["wrong_penalty"], f"COLLECT_FAIL: no pending target at {aisle}"373 374        if self.cfg["closest_rule"]:375            req_aisle = self._target_products[self._closest_idx]["aisle"]376            if aisle != req_aisle:377                d_here = self._manhattan(AISLE_GRID_POS[aisle])378                d_req  = self._manhattan(AISLE_GRID_POS[req_aisle])379                return (380                    self.cfg["wrong_penalty"],381                    f"CLOSEST_RULE_FAIL: went to {aisle}({d_here} steps) "382                    f"but {req_aisle}({d_req} steps) is closest",383                )384 385        idx = candidates[0]386        self._collected[idx] = True387        item_reward = self._target_products[idx]["reward"] * self.cfg["reward_mult"]388        return item_reward, f"COLLECTED {self._target_products[idx]['name']} +{item_reward:.1f}"389 390    def _build_obs(self) -> np.ndarray:391        obs = np.zeros(OBS_LEN, dtype=np.float32)392        obs[0] = self._agent_pos[0] / (self.GRID_ROWS - 1)393        obs[1] = self._agent_pos[1] / (self.GRID_COLS - 1)394        obs[2] = float(self._phase)395        obs[3] = self._steps_taken / self.cfg["max_steps"]396        for i in range(MAX_PRODUCTS):397            base = 4 + i * 3398            if i < len(self._target_products):399                p = self._target_products[i]400                gr, gc = AISLE_GRID_POS[p["aisle"]]401                obs[base]     = gr / (self.GRID_ROWS - 1)402                obs[base + 1] = gc / (self.GRID_COLS - 1)403                obs[base + 2] = float(self._collected[i])404        return obs405 406    def _build_info(self, event: str, step_reward: float) -> Dict[str, Any]:407        return {408            "task_status":      self._task_status,409            "reward_log":       self._reward_log,410            "total_reward":     round(self._total_reward, 4),411            "normalised_score": self.normalised_score(),412            "step_reward":      round(step_reward, 4),413            "steps_taken":      self._steps_taken,414            "steps_remaining":  self.cfg["max_steps"] - self._steps_taken,415            "inventory":        self.collected_names,416            "targets":          self.target_names,417            "level":            self.level,418            "phase":            self._phase,419            "phase_label":      "collecting" if self._phase == 0 else "returning_to_counter",420            "event":            event,421            "action_mask":      self.action_masks().tolist(),422            "agent_pos":        list(self._agent_pos),423            "closest_target":   (self._target_products[self._closest_idx]["name"]424                                 if not all(self._collected) else None),425            "closest_rule":     self.cfg["closest_rule"],426            "obs_len":          OBS_LEN,427        }428 429    def _append_log(self, action: int, reward: float, event: str):430        self._reward_log.append({431            "step":         self._steps_taken,432            "action":       ACTION_NAMES[action],433            "step_reward":  round(reward, 4),434            "total_reward": round(self._total_reward, 4),435            "position":     list(self._agent_pos),436            "phase":        self._phase,437            "inventory":    self.collected_names.copy(),438            "task_status":  self._task_status,439            "event":        event,440        })441 442    def _resolve_products(self) -> List[Dict]:443        n = self.cfg["num_products"]444        if self._fixed_products:445            out = []446            for name in self._fixed_products[:n]:447                key = name.lower()448                if key not in ALL_PRODUCTS:449                    raise ValueError(f"Unknown product '{name}'")450                out.append(ALL_PRODUCTS[key])451            return out452        return random.sample(list(ALL_PRODUCTS.values()), min(n, len(ALL_PRODUCTS)))453 454 455 456if GYM_AVAILABLE:457    for _lvl in ("easy", "medium", "hard"):458        try:459            gym.register(460                f"SupermarketNav-{_lvl.capitalize()}-v0",461                entry_point="supermart_env:SupermarketEnv",462                kwargs={"level": _lvl},463            )464        except Exception:465            pass