Spaces:
Running
Running
| """Real, paused ViZDoom episodes using the bundled Freedoom 2 assets. | |
| The action sets, durations, and model telemetry match the recorded JEV demo. | |
| This module has no policy: the caller must explicitly select every action. | |
| """ | |
| from __future__ import annotations | |
| import base64 | |
| import io | |
| import json | |
| import secrets | |
| from pathlib import Path | |
| import vizdoom as vzd | |
| from PIL import Image | |
| ENEMIES = { | |
| "Cacodemon", "DoomImp", "Demon", "Spectre", "HellKnight", "BaronOfHell", | |
| "Arachnotron", "Revenant", "Mancubus", "ZombieMan", "ShotgunGuy", | |
| "ChaingunGuy", "LostSoul", "PainElemental", "Cyberdemon", "SpiderMastermind", | |
| } | |
| def action_set(verb: str, direction: str) -> list[dict]: | |
| return [ | |
| {"id": f"{verb}_left_fine", "label": "Left · fine", "buttons": [1, 0, 0], | |
| "tics": 2, "description": f"{direction} left for 2 tics, a small alignment adjustment"}, | |
| {"id": f"{verb}_left", "label": "Left", "buttons": [1, 0, 0], | |
| "tics": 8, "description": f"{direction} left for 8 tics"}, | |
| {"id": f"{verb}_right_fine", "label": "Right · fine", "buttons": [0, 1, 0], | |
| "tics": 2, "description": f"{direction} right for 2 tics, a small alignment adjustment"}, | |
| {"id": f"{verb}_right", "label": "Right", "buttons": [0, 1, 0], | |
| "tics": 8, "description": f"{direction} right for 8 tics"}, | |
| {"id": "attack", "label": "Attack", "buttons": [0, 0, 1], | |
| "tics": 12, "description": "Fire weapon for 12 tics at current crosshair direction"}, | |
| {"id": "wait", "label": "Wait", "buttons": [0, 0, 0], | |
| "tics": 4, "description": "Wait without moving for 4 tics"}, | |
| ] | |
| SCENARIOS = { | |
| "basic": { | |
| "instructions": "You are playing Doom target practice. One stationary visible monster across the room must be shot. You can strafe LEFT/RIGHT (no turning). Crosshair is x=320, screen width 640. Moving LEFT shifts objects RIGHT on screen; moving RIGHT shifts objects LEFT. Align horizontal aim to monster then ATTACK; Doom auto-aims vertically. One hit kills. More score = fewer elapsed tics and missed shots.", | |
| "actions": action_set("move", "Strafe"), | |
| }, | |
| "defend_the_center": { | |
| "instructions": "You are playing Doom defend-the-center. Survive and kill approaching monsters. You stand still and can TURN LEFT/RIGHT or ATTACK. Crosshair is x=320, screen width 640. Turning LEFT shifts objects RIGHT on screen; turning RIGHT shifts objects LEFT. Aim horizontally at a visible monster and fire; Doom auto-aims vertically. If no visible living monster, turn to search. Each monster initially dies in one hit. Use fine turns when almost aligned. Dead bodies are not targets.", | |
| "actions": action_set("turn", "Turn"), | |
| }, | |
| } | |
| class DoomEpisode: | |
| """All methods for an instance are called on one dedicated worker thread.""" | |
| def __init__(self, scenario: str): | |
| self.scenario = scenario | |
| self.spec = SCENARIOS[scenario] | |
| self.steps = 0 | |
| self.previous: list[dict] = [] | |
| self.screen = None | |
| self.last_telemetry: dict = {} | |
| self.game = vzd.DoomGame() | |
| try: | |
| self.game.load_config(str(Path(vzd.scenarios_path) / f"{scenario}.cfg")) | |
| self.game.set_doom_game_path(str(Path(vzd.__file__).parent / "freedoom2.wad")) | |
| self.game.set_window_visible(False) | |
| self.game.set_sound_enabled(False) | |
| self.game.set_labels_buffer_enabled(True) | |
| self.game.set_screen_format(vzd.ScreenFormat.RGB24) | |
| self.game.set_screen_resolution(vzd.ScreenResolution.RES_640X480) | |
| self.game.set_render_crosshair(True) | |
| self.game.set_seed(secrets.randbelow(2**31 - 1)) | |
| self.game.set_mode(vzd.Mode.PLAYER) | |
| self.game.init() | |
| self.game.new_episode() | |
| self._capture() | |
| except Exception: | |
| self.game.close() | |
| raise | |
| def _capture(self) -> None: | |
| state = self.game.get_state() | |
| if state is None: | |
| return | |
| self.screen = state.screen_buffer.copy() | |
| labels = [ | |
| { | |
| "name": obj.object_name, | |
| "screen_left": int(obj.x), | |
| "screen_right": int(obj.x + obj.width), | |
| "screen_center_x": round(obj.x + obj.width / 2, 1), | |
| "screen_width": int(obj.width), | |
| "screen_top": int(obj.y), | |
| "screen_height": int(obj.height), | |
| } | |
| for obj in state.labels if obj.object_name in ENEMIES | |
| ] | |
| self.last_telemetry = { | |
| "tic": int(self.game.get_episode_time()), "crosshair_x": 320, | |
| "screen_width": 640, "visible_monsters": labels, | |
| "game_variables": { | |
| var.name: float(self.game.get_game_variable(var)) | |
| for var in [vzd.GameVariable.HEALTH, vzd.GameVariable.AMMO2, | |
| vzd.GameVariable.KILLCOUNT, vzd.GameVariable.POSITION_X, | |
| vzd.GameVariable.POSITION_Y, vzd.GameVariable.ANGLE] | |
| }, | |
| } | |
| def snapshot(self) -> dict: | |
| self._capture() | |
| image = io.BytesIO() | |
| Image.fromarray(self.screen).save(image, format="PNG", compress_level=3) | |
| return { | |
| "scenario": self.scenario, | |
| "frame": base64.b64encode(image.getvalue()).decode("ascii"), | |
| "metrics": { | |
| "health": float(self.game.get_game_variable(vzd.GameVariable.HEALTH)), | |
| "ammo": float(self.game.get_game_variable(vzd.GameVariable.AMMO2)), | |
| "kills": float(self.game.get_game_variable(vzd.GameVariable.KILLCOUNT)), | |
| "reward": float(self.game.get_total_reward()), "steps": self.steps, | |
| }, | |
| "done": bool(self.game.is_episode_finished()), | |
| "actions": [{key: value for key, value in action.items() if key != "buttons"} | |
| for action in self.spec["actions"]], | |
| "telemetry": self.last_telemetry, | |
| } | |
| def decision_payload(self) -> dict: | |
| self._capture() | |
| return { | |
| "kind": "choice", | |
| "state": self.spec["instructions"] + "\nCurrent engine telemetry:\n" | |
| + json.dumps(self.last_telemetry, separators=(",", ":")) | |
| + "\nRecent decisions:\n" | |
| + json.dumps(self.previous[-3:], separators=(",", ":")), | |
| "question": "Choose one action now to maximize the game objective.", | |
| "options": [a["description"] for a in self.spec["actions"]], | |
| } | |
| def step(self, action_id: str) -> dict: | |
| action = next(a for a in self.spec["actions"] if a["id"] == action_id) | |
| if self.game.is_episode_finished(): | |
| raise ValueError("The episode has finished. Start a new game.") | |
| before = self.last_telemetry | |
| reward = 0.0 | |
| # Capture each actual tic so the final response retains the last engine frame. | |
| # PLAYER mode advances only here, never while awaiting network inference. | |
| for _ in range(action["tics"]): | |
| if self.game.is_episode_finished(): | |
| break | |
| reward += self.game.make_action(action["buttons"], 1) | |
| self._capture() | |
| self.steps += 1 | |
| self.previous.append({ | |
| "action": action["description"], | |
| "before_targets": before.get("visible_monsters", []), | |
| "after_targets": self.last_telemetry.get("visible_monsters", []), | |
| "reward": reward, | |
| }) | |
| self.previous = self.previous[-3:] | |
| return self.snapshot() | |
| def close(self) -> None: | |
| self.game.close() | |