Spaces:
Running
Running
| # Copyright (c) Meta Platforms, Inc. and affiliates. | |
| # All rights reserved. | |
| # | |
| # This source code is licensed under the BSD-style license found in the | |
| # LICENSE file in the root directory of this source tree | |
| import uuid | |
| from typing import Any, Dict, List, Optional, Tuple | |
| try: | |
| from openenv.core.env_server import Environment | |
| from openenv.core.env_server.types import State | |
| except ImportError: | |
| from core.env_server import Environment | |
| from core.env_server.types import State | |
| from ..models import GridWorldAction, GridWorldObservation, MoveAction | |
| class GridWorldEnvironment(Environment): | |
| """ | |
| A simple 5x5 Grid World environment. | |
| The agent starts at [0, 0] and must navigate to [4, 4]. | |
| """ | |
| SUPPORTS_CONCURRENT_SESSIONS = True | |
| def __init__(self): | |
| super().__init__() | |
| # ---Environment Configuration --- | |
| self.grid_size = 5 | |
| self.goal_pos = [4, 4] | |
| # --- Internal State Variable --- | |
| self.agent_x = 0 | |
| self.agent_y = 0 | |
| # Initialize the base OpenEnv State Container | |
| self._state = State(episode_id=str(uuid.uuid4()), step_count=0) | |
| def reset(self) -> GridWorldObservation: | |
| # Update State | |
| self.agent_x = 0 | |
| self.agent_y = 0 | |
| self._state.step_count = 0 | |
| self._state.episode_id = str(uuid.uuid4()) | |
| # Return initial observation (reward must be float 0.0, not None) | |
| return GridWorldObservation( | |
| x=self.agent_x, | |
| y=self.agent_y, | |
| message="Welcome to Grid World! Goal is at [4, 4].", | |
| reward=0.0, | |
| done=False, | |
| ) | |
| def step(self, action: GridWorldAction) -> GridWorldObservation: | |
| # Increment step counter in the base State | |
| self._state.step_count += 1 | |
| # ============================================= | |
| move = action.action | |
| # self._state.episode_steps += 1 | |
| # Use current state | |
| # current_x = self._state.agent_x | |
| # current_y = self._state.agent_y | |
| move = action.action | |
| if move == MoveAction.UP: | |
| self.agent_x -= 1 | |
| elif move == MoveAction.DOWN: | |
| self.agent_x += 1 | |
| elif move == MoveAction.LEFT: | |
| self.agent_y -= 1 | |
| elif move == MoveAction.RIGHT: | |
| self.agent_y += 1 | |
| # Clamp to boundaries | |
| self.agent_x = max(0, min(self.agent_x, self.grid_size - 1)) | |
| self.agent_y = max(0, min(self.agent_y, self.grid_size - 1)) | |
| # # Update State | |
| # self._state.agent_x = current_x | |
| # self._state.agent_y = current_y | |
| # Logic | |
| done = False | |
| message = "Keep going..." | |
| reward = -0.1 | |
| if [self.agent_x, self.agent_y] == self.goal_pos: | |
| reward = 1.0 | |
| done = True | |
| message = "You found the goal!" | |
| return GridWorldObservation( | |
| x=self.agent_x, y=self.agent_y, message=message, reward=reward, done=done | |
| ) | |
| def web_actions( | |
| self, observation: Dict[str, Any] | |
| ) -> List[Tuple[str, Dict[str, Any]]]: | |
| """The four moves as buttons.""" | |
| arrows = {"UP": "↑", "DOWN": "↓", "LEFT": "←", "RIGHT": "→"} | |
| return [ | |
| (f"{arrows[m.value]} {m.value.lower()}", {"action": m.value}) | |
| for m in MoveAction | |
| ] | |
| def render_web(self, observation: Dict[str, Any]) -> Optional[str]: | |
| """Draw the grid with the agent (a dot) and the goal (a star).""" | |
| if "x" not in observation: | |
| return None | |
| cells = [] | |
| for row in range(self.grid_size): | |
| for col in range(self.grid_size): | |
| agent = [row, col] == [observation["x"], observation["y"]] | |
| goal = [row, col] == self.goal_pos | |
| border = "var(--color-accent)" if goal else "transparent" | |
| dot = "" | |
| if agent: | |
| dot = '<span style="width:60%;height:60%;border-radius:50%;background:var(--color-accent)"></span>' | |
| elif goal: | |
| dot = "★" | |
| cells.append( | |
| '<span style="display:flex;align-items:center;justify-content:center;' | |
| "background:var(--border-color-primary);border-radius:6px;" | |
| f'border:2px solid {border};color:var(--color-accent);font-size:20px">{dot}</span>' | |
| ) | |
| return ( | |
| '<div role="img" aria-label="Grid World board" style="display:inline-grid;' | |
| f"grid-template-columns:repeat({self.grid_size},40px);grid-auto-rows:40px;gap:4px;padding:10px;" | |
| 'border:1px solid var(--border-color-primary);border-radius:10px;background:var(--background-fill-secondary)">' | |
| + "".join(cells) | |
| + "</div>" | |
| ) | |
| def state(self) -> State: | |
| return self._state | |