범용 JS/AI 게임 제작 가이드 (v3.6)

import gymnasium as gym
from gymnasium import spaces
import numpy as np
import asyncio
import websockets
import json
import logging
import threading
import queue

# Websockets 라이브러리의 상세 로그 비활성화
logging.getLogger("websockets").setLevel(logging.ERROR)
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(message)s')

class ServerThread(threading.Thread):
    """
    백그라운드에서 웹소켓 서버를 실행하고 메인 스레드와 큐를 통해 통신하는
    독립적인 스레드입니다.
    """
    def __init__(self):
        super().__init__()
        self.loop = asyncio.new_event_loop()
        self.command_queue = queue.Queue()
        self.observation_queue = queue.Queue()
        self.daemon = True # 메인 스레드가 종료되면 함께 종료

    # --- FIX ---
    # 최신 websockets 라이브러리(v10.0+)는 핸들러에 'path' 인자를 전달하지 않습니다.
    # 따라서 함수 정의에서 'path'를 제거해야 합니다.
    async def handler(self, websocket):
        logging.info(f"새로운 클라이언트 연결됨: {websocket.remote_address}")
        consumer_task = asyncio.ensure_future(self.consumer(websocket))
        producer_task = asyncio.ensure_future(self.producer(websocket))
        done, pending = await asyncio.wait(
            [consumer_task, producer_task],
            return_when=asyncio.FIRST_COMPLETED,
        )
        for task in pending:
            task.cancel()
        logging.info(f"클라이언트 연결 종료: {websocket.remote_address}")

    async def consumer(self, websocket):
        """메인 스레드로부터 명령을 받아 클라이언트로 전송"""
        while True:
            try:
                command = await self.loop.run_in_executor(None, self.command_queue.get)
                await websocket.send(json.dumps(command))
            except (websockets.exceptions.ConnectionClosed, asyncio.CancelledError):
                break

    async def producer(self, websocket):
        """클라이언트로부터 관측 데이터를 받아 메인 스레드로 전송"""
        while True:
            try:
                message = await websocket.recv()
                data = json.loads(message)
                if "observation" in data:
                    obs = np.array(data["observation"], dtype=np.float64)
                    reward = data.get("reward", 0)
                    done = data.get("done", False)
                    self.observation_queue.put((obs, reward, done))
            except (websockets.exceptions.ConnectionClosed, asyncio.CancelledError):
                break

    def run(self):
        asyncio.set_event_loop(self.loop)
        async def start_server():
            async with websockets.serve(self.handler, "localhost", 8765):
                await asyncio.Future()
        try:
            # run_until_complete() 대신 run()을 사용하여 서버를 실행합니다.
            # 이 방식은 스레드 종료 시 더 깔끔하게 처리됩니다.
            self.loop.run_until_complete(start_server())
        finally:
            self.loop.close()

    def stop(self):
        if self.loop.is_running():
            self.loop.call_soon_threadsafe(self.loop.stop)
            logging.info("웹소켓 서버 종료를 시도합니다.")

class WebSocketEnv(gym.Env):
    def __init__(self, observation_shape, action_space_config):
        super(WebSocketEnv, self).__init__()
        self.observation_space = spaces.Box(low=-np.inf, high=np.inf, shape=observation_shape, dtype=np.float64)
        if action_space_config['type'] == 'discrete': self.action_space = spaces.Discrete(action_space_config['n'])
        elif action_space_config['type'] == 'continuous': self.action_space = spaces.Box(low=np.array(action_space_config['low']), high=np.array(action_space_config['high']), dtype=np.float64)
        self.server_thread = ServerThread()
        self.server_thread.start()
        print("첫 클라이언트의 연결 및 데이터 수신을 기다립니다...")

    def reset(self, seed=None, options=None):
        super().reset(seed=seed)
        with self.server_thread.observation_queue.mutex:
            self.server_thread.observation_queue.queue.clear()
        self.server_thread.command_queue.put({"command": "reset"})
        obs, _, _ = self.server_thread.observation_queue.get()
        return obs, {}

    def step(self, action):
        action_to_send = action.tolist() if isinstance(action, np.ndarray) else float(action)
        self.server_thread.command_queue.put({"command": "action", "action": action_to_send})
        obs, reward, done = self.server_thread.observation_queue.get()
        return obs, reward, done, False, {}

    def render(self, mode='human'): pass
   
    def close(self):
        self.server_thread.stop()
        print("웹소켓 서버가 종료되었습니다.")

import torch
import torch.nn as nn
from stable_baselines3 import PPO
import numpy as np
import os

# --- 설정 ---
MODEL_DIR = "model"  # 모델이 저장된 폴더 경로
MODEL_FILENAME = "final_model.zip"
MODEL_PATH = os.path.join(MODEL_DIR, MODEL_FILENAME)

# 변환된 ONNX 모델을 저장할 경로와 파일명
# index.html 파일이 있는 루트 폴더에 저장해야 AI 모드가 인식할 수 있습니다.
ONNX_OUTPUT_PATH = "model.onnx"

# train.py에 정의된 환경의 관찰 공간 모양
# (배치 크기 1을 추가하여 [1, 20] 형태로 만들어야 합니다)
OBSERVATION_SHAPE = (1, 20) 

def convert_to_onnx():
    """
    Stable-Baselines3 PPO 모델(.zip)을 ONNX 형식으로 변환합니다.
    """
    print(f"'{MODEL_PATH}'에서 학습된 모델을 불러옵니다...")

    # 모델 파일이 존재하는지 확인
    if not os.path.exists(MODEL_PATH):
        print(f"오류: 모델 파일 '{MODEL_PATH}'을 찾을 수 없습니다.")
        print("먼저 train.py를 실행하여 모델을 학습시키고 저장해야 합니다.")
        return

    # Stable-Baselines3 모델 불러오기
    # PyTorch를 사용하도록 device='cpu'로 명시하는 것이 안정적입니다.
    try:
        model = PPO.load(MODEL_PATH, device='cpu')
    except Exception as e:
        print(f"모델을 불러오는 중 오류가 발생했습니다: {e}")
        return
        
    print("모델을 성공적으로 불러왔습니다.")

    # ONNX 변환을 위해 모델을 평가 모드(evaluation mode)로 설정합니다.
    # 이 과정은 드롭아웃(dropout)이나 배치 정규화(batch normalization) 등의
    # 학습용 레이어를 비활성화하여 일관된 출력을 보장합니다.
    model.policy.eval()

    # ONNX 변환을 위한 더미 입력 데이터를 생성합니다.
    # 이 데이터는 모델의 구조를 추적하는 데 사용되며, 실제 입력값은 중요하지 않습니다.
    dummy_input = torch.randn(OBSERVATION_SHAPE, dtype=torch.float32)

    print(f"모델을 ONNX 형식으로 변환합니다. 저장 위치: '{ONNX_OUTPUT_PATH}'")

    try:
        # torch.onnx.export 함수를 사용하여 모델을 변환합니다.
        torch.onnx.export(
            model.policy,                 # 변환할 모델의 정책 네트워크
            dummy_input,                  # 모델 구조 추적을 위한 더미 입력
            ONNX_OUTPUT_PATH,             # 저장될 ONNX 파일 경로
            opset_version=11,             # ONNX 버전 (11은 호환성이 좋음)
            input_names=['input'],        # ONNX 모델의 입력 레이어 이름
            output_names=['output'],      # ONNX 모델의 출력 레이어 이름
            dynamic_axes={'input': {0: 'batch_size'}, # 배치 크기가 변할 수 있도록 설정
                          'output': {0: 'batch_size'}}
        )
        print("-" * 50)
        print(" ONNX 변환이 성공적으로 완료되었습니다!")
        print(f"이제 'index.html' 파일을 브라우저에서 열고 'AI 모드'를 테스트할 수 있습니다.")
        print("-" * 50)

    except Exception as e:
        print(f"ONNX 변환 중 오류가 발생했습니다: {e}")

if __name__ == '__main__':
    convert_to_onnx()