예제 소스
OpenAI Gym과 강화학습(Reinforcement Learning) · 파일 12개
Pong-vlm.py
#!/usr/bin/env python3
"""
Pong-v4 를 로컬 VLM(Qwen3-VL 등)으로 플레이합니다.
- `--vlm-input image`(기본): 최근 N장 RGB 를 버퍼에 두고, **시간 순(과거 → 최근)** 으로 모델에 전달합니다.
- `--vlm-input text`: 화면 대신 **추정 볼 좌표·패들 y·motion 등 로그 형태 히스토리**만 LLM에 넣어 액션을 받습니다 (느린 비전 호출 회피).
- 기본 백엔드: `transformers` 의 `image-text-to-text` 파이프라인 (로컬 `models/` 경로).
- 선택 백엔드: OpenAI vision API (`--backend openai`).
모델 디렉터리
------------
동일 볼륨에서 `cp -al` 로 복사해 두었습니다 (inode 공유, 용량 중복 없음):
OpenAiGym/models/Qwen__Qwen3-VL-2B-Instruct
없으면:
pip install huggingface_hub torch torchvision transformers accelerate sentencepiece pillow
python download_vlm_to_models.py
실행 예
-------
python Pong-vlm.py --device mps --frames 6 --stride 4 --episodes 1 --no-gui
python Pong-vlm.py --vlm-input text --stride 2 --text-history 20 --device mps --episodes 1
→ 볼/패들 수치 로그만 모델에 전달 (이미지 없음).
python Pong-vlm.py --assist intercept --device mps --frames 4 --stride 8
→ VLM은 느리게 갱신해도 매 프레임 프레임차 기반으로 패들 y를 공에 맞춤.
python Pong-vlm.py --print-track --print-track-every 2 --no-gui
→ 공 추정(ball_cx/cy)·패들 y·motion_sum 등을 stdout에 반복 출력(assist 없이도 가능).
"""
from __future__ import annotations
import argparse
import io
import json
import os
import re
import sys
import urllib.error
import urllib.request
from collections import deque
from typing import Any, Deque, List, Optional
import ale_py
import gymnasium as gym
import numpy as np
from PIL import Image
gym.register_envs(ale_py)
from PyQt5.QtCore import Qt
from PyQt5.QtGui import QImage, QPixmap
from PyQt5.QtWidgets import QApplication, QLabel, QVBoxLayout, QWidget
_ROOT = os.path.dirname(os.path.abspath(__file__))
_DEFAULT_LOCAL = os.path.join(_ROOT, "models", "Qwen__Qwen3-VL-2B-Instruct")
class AtariQtView(QWidget):
def __init__(self, title: str = "Pong VLM"):
super().__init__()
self.setWindowTitle(title)
self._label = QLabel(self)
self._label.setAlignment(Qt.AlignCenter)
lay = QVBoxLayout(self)
lay.setContentsMargins(0, 0, 0, 0)
lay.addWidget(self._label)
def show_frame(self, rgb: np.ndarray) -> None:
rgb = np.ascontiguousarray(np.asarray(rgb, dtype=np.uint8))
if rgb.ndim != 3 or rgb.shape[2] != 3:
return
h, w, _ = rgb.shape
qimg = QImage(rgb.data, w, h, 3 * w, QImage.Format_RGB888).copy()
self._label.setPixmap(QPixmap.fromImage(qimg))
app = QApplication.instance()
if app is not None:
app.processEvents()
def _rgb_to_pil(rgb: np.ndarray, max_side: int) -> Image.Image:
im = Image.fromarray(np.asarray(rgb, dtype=np.uint8), mode="RGB")
w, h = im.size
scale = min(1.0, float(max_side) / max(w, h))
if scale < 1.0:
im = im.resize((max(1, int(w * scale)), max(1, int(h * scale))), Image.Resampling.BILINEAR)
return im
def _pong_prompt_text(score: float, last_reward: float) -> str:
return "\n".join(
[
"당신은 Atari Pong 의 오른쪽 패들을 조작합니다.",
"이미지들은 시간 순서(첫 장이 가장 과거, 마지막이 가장 최근)입니다.",
"공의 방향·속도 변화를 보고 패들 위치와, 맞을 때 강/약 타격(속도)까지 고려하세요.",
"",
"액션 번호 (0~5 정수 하나만 JSON):",
"0=NOOP, 1=FIRE, 2=RIGHT, 3=LEFT, 4=RIGHTFIRE, 5=LEFTFIRE",
"(오른쪽 패들) 2=RIGHT → 패들이 화면 위쪽으로, 3=LEFT → 패들이 화면 아래쪽으로 움직입니다.",
"",
f"에피소드 누적 점수: {score:.1f}, 직전 step 보상: {last_reward:.3f}",
"",
'출력은 한 줄 JSON 만: {"action": 2, "reason": "한글 한 문장"}',
]
)
def _pong_state_prompt(score: float, last_reward: float, csv_block: str, h: int, w: int) -> str:
return "\n".join(
[
"당신은 Atari Pong 의 오른쪽 패들을 조작합니다.",
"아래 표는 **프레임 차이 기반 추정** 값입니다 (실제 픽셀과 약간 다를 수 있음).",
"첫 데이터 행이 가장 과거, 마지막 행이 가장 최근입니다.",
f"화면 크기: 높이 {h}, 너비 {w} (y는 위가 작음).",
"",
"CSV 컬럼:",
"step,ball_cx,ball_cy,ball_cy_pred,paddle_py,motion_sum,delta_y,reject,ghost,step_reward",
"",
csv_block,
"",
f"에피소드 누적 점수: {score:.1f}, 직전 step 보상(위 표 마지막 행의 step_reward): {last_reward:.3f}",
"",
"액션 번호 (0~5 정수 하나만 JSON):",
"0=NOOP, 1=FIRE, 2=RIGHT, 3=LEFT, 4=RIGHTFIRE, 5=LEFTFIRE",
"(오른쪽 패들) 2=RIGHT → 패들이 화면 위쪽으로, 3=LEFT → 패들이 화면 아래쪽으로 움직입니다.",
"",
'출력은 한 줄 JSON 만: {"action": 2, "reason": "한글 한 문장"}',
]
)
def _telemetry_csv_row(step: int, t: dict[str, Any], step_reward: float) -> str:
def cell(v: Any) -> str:
if v is None:
return ""
if isinstance(v, float):
return f"{v:.2f}"
return str(v)
return ",".join(
[
str(int(step)),
cell(t.get("ball_cx")),
cell(t.get("ball_cy")),
cell(t.get("ball_cy_pred")),
cell(t.get("paddle_py")),
cell(t.get("motion_sum")),
cell(t.get("delta_y")),
str(t.get("reject") or ""),
str(t.get("ghost") if t.get("ghost") is not None else ""),
f"{step_reward:.3f}",
]
)
class StateHistoryBuffer:
"""텍스트 모드용: 최근 스텝의 볼·패들 추정치를 CSV 한 줄씩 보관."""
def __init__(self, maxlen: int) -> None:
self._rows: Deque[str] = deque(maxlen=max(2, maxlen))
def clear(self) -> None:
self._rows.clear()
def push(self, step: int, telemetry: dict[str, Any], step_reward: float) -> None:
self._rows.append(_telemetry_csv_row(step, telemetry, step_reward))
def ready(self) -> bool:
return len(self._rows) >= 1
def csv_block(self) -> str:
return "\n".join(self._rows)
def _parse_action_json(text: str) -> int:
text = text.strip()
m = re.search(r'\{[^{}]*"action"[^{}]*\}', text, flags=re.I | re.DOTALL)
if m:
try:
obj = json.loads(m.group(0))
a = int(obj["action"])
if 0 <= a <= 5:
return a
except (json.JSONDecodeError, KeyError, TypeError, ValueError):
pass
m2 = re.search(r"\b([0-5])\b", text)
if m2:
return int(m2.group(1))
return 0
def _extract_generated_text(raw: Any) -> str:
if isinstance(raw, list) and raw:
first = raw[0]
if isinstance(first, dict):
gen = first.get("generated_text", "")
if isinstance(gen, list):
for msg in reversed(gen):
if isinstance(msg, dict) and msg.get("role") == "assistant":
content = msg.get("content", "")
if isinstance(content, list):
parts = []
for p in content:
if isinstance(p, dict) and "text" in p:
parts.append(str(p["text"]))
return "\n".join(parts)
return str(content)
return str(gen)
return str(gen)
return str(first)
return str(raw)
def _resolve_pipeline_device(name: str):
name = (name or "auto").lower()
if name == "cpu":
return -1
if name in ("cuda", "gpu"):
try:
import torch
if torch.cuda.is_available():
return 0
except Exception:
pass
return -1
if name in ("mps", "metal"):
try:
import torch
if torch.backends.mps.is_available():
return "mps"
except Exception:
pass
return -1
if name == "auto":
try:
import torch
if torch.cuda.is_available():
return 0
if torch.backends.mps.is_available():
return "mps"
except Exception:
pass
return -1
return -1
class LocalVlmPong:
def __init__(self, model_path: str, device: str) -> None:
self._model_path = model_path
self._device = _resolve_pipeline_device(device)
self._pipe = None
def _pipe_lazy(self):
if self._pipe is not None:
return self._pipe
if not os.path.isdir(self._model_path):
raise FileNotFoundError(
f"모델 경로 없음: {self._model_path}\n"
" python download_vlm_to_models.py\n"
" 또는 rl 쪽 가중치를 cp -al 로 models/ 에 두세요."
)
try:
from transformers import pipeline
except ImportError as e:
raise RuntimeError("pip install transformers torch accelerate sentencepiece") from e
self._pipe = pipeline(
"image-text-to-text",
model=self._model_path,
trust_remote_code=True,
device=self._device,
)
return self._pipe
def _gen_kwargs(self, pipe) -> dict:
"""모델 기본 generation_config 의 max_length 와 max_new_tokens 충돌 경고를 피함."""
tok = getattr(pipe, "tokenizer", None)
pad = getattr(tok, "eos_token_id", None) if tok is not None else None
try:
from transformers import GenerationConfig
except ImportError:
g: dict = {"max_new_tokens": 48, "do_sample": False}
if pad is not None:
g["pad_token_id"] = pad
return g
base: dict = {}
m = getattr(pipe, "model", None)
gc = getattr(m, "generation_config", None) if m is not None else None
if gc is not None and hasattr(gc, "to_dict"):
base = dict(gc.to_dict())
for k in ("max_length", "temperature", "top_p", "top_k"):
base.pop(k, None)
base["max_new_tokens"] = 48
base["do_sample"] = False
if pad is not None:
base["pad_token_id"] = pad
try:
cfg = GenerationConfig.from_dict(base)
return {"generation_config": cfg}
except Exception:
g = {"max_new_tokens": 48, "do_sample": False}
if pad is not None:
g["pad_token_id"] = pad
return g
def predict(self, rgb_frames: List[np.ndarray], max_side: int, score: float, last_r: float) -> int:
pipe = self._pipe_lazy()
pils = [_rgb_to_pil(f, max_side) for f in rgb_frames]
prompt = _pong_prompt_text(score, last_r)
content: list = []
for i, pil in enumerate(pils):
content.append({"type": "text", "text": f"[frame {i + 1}/{len(pils)}]"})
content.append({"type": "image", "image": pil})
content.append({"type": "text", "text": prompt})
messages = [{"role": "user", "content": content}]
gen = self._gen_kwargs(pipe)
raw = None
last_err: Optional[BaseException] = None
for caller in (
lambda m=messages: pipe(text=m, generate_kwargs=gen),
lambda: pipe(images=pils[-1], text=prompt, generate_kwargs=gen),
):
try:
raw = caller()
break
except Exception as e:
last_err = e
continue
if raw is None:
raise RuntimeError(f"VLM 추론 실패: {last_err!r}") from last_err
text = _extract_generated_text(raw)
return _parse_action_json(text)
def predict_text(self, user_prompt: str) -> int:
"""이미지 없이 텍스트만 입력 (상태 로그 모드)."""
pipe = self._pipe_lazy()
content: list = [{"type": "text", "text": user_prompt}]
messages = [{"role": "user", "content": content}]
gen = self._gen_kwargs(pipe)
raw = None
last_err: Optional[BaseException] = None
for caller in (
lambda m=messages: pipe(text=m, generate_kwargs=gen),
lambda: pipe(text=user_prompt, generate_kwargs=gen),
):
try:
raw = caller()
break
except Exception as e:
last_err = e
continue
if raw is None:
raise RuntimeError(f"VLM 텍스트 추론 실패: {last_err!r}") from last_err
text = _extract_generated_text(raw)
return _parse_action_json(text)
def _resize_jpeg_b64(rgb: np.ndarray, max_side: int, quality: int) -> str:
import base64
im = _rgb_to_pil(rgb, max_side)
buf = io.BytesIO()
im.save(buf, format="JPEG", quality=quality)
return base64.standard_b64encode(buf.getvalue()).decode("ascii")
def _build_openai_messages(frames_b64: List[str], score: float, last_reward: float) -> list:
text = _pong_prompt_text(score, last_reward)
parts: list = [{"type": "text", "text": text}]
for i, b64 in enumerate(frames_b64):
parts.append({"type": "text", "text": f"[frame {i + 1}/{len(frames_b64)}]"})
parts.append(
{"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{b64}", "detail": "low"}}
)
return [
{"role": "system", "content": "You are a concise game controller. Output only valid JSON."},
{"role": "user", "content": parts},
]
def _build_openai_messages_text(state_user_prompt: str) -> list:
return [
{"role": "system", "content": "You are a concise game controller. Output only valid JSON."},
{"role": "user", "content": state_user_prompt},
]
def _call_openai_chat(api_key: str, base_url: str, model: str, messages: list) -> str:
url = base_url.rstrip("/") + "/chat/completions"
body = json.dumps(
{"model": model, "messages": messages, "max_tokens": 120, "temperature": 0.2}
).encode("utf-8")
req = urllib.request.Request(
url,
data=body,
headers={"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"},
method="POST",
)
try:
with urllib.request.urlopen(req, timeout=120) as resp:
payload = json.loads(resp.read().decode("utf-8"))
except urllib.error.HTTPError as e:
err = e.read().decode("utf-8", errors="replace")
raise RuntimeError(f"VLM HTTP {e.code}: {err}") from None
return payload["choices"][0]["message"]["content"]
class FrameBuffer:
def __init__(self, maxlen: int) -> None:
self._buf: Deque[np.ndarray] = deque(maxlen=maxlen)
def push(self, rgb: np.ndarray) -> None:
self._buf.append(np.ascontiguousarray(np.asarray(rgb, dtype=np.uint8)))
def ready(self) -> bool:
return len(self._buf) >= 1
def as_list_oldest_first(self) -> List[np.ndarray]:
return list(self._buf)
class InterceptAssist:
"""
프레임 차이로 공의 대략적인 (x,y)를 잡고, 오른쪽 패들의 세로 중심과 맞추는 휴리스틱.
ALE Pong-v4 오른쪽 패들: 2=RIGHT(화면 위로, y 감소), 3=LEFT(화면 아래로, y 증가).
공이 패들 근처에서 차분이 작아지면 궤적이 끊기므로, ROI를 오른쪽까지 넓히고
마지막으로 잡힌 속도로 짧게 coast 예측(--assist-ghost-max)합니다.
"""
def __init__(
self,
deadzone: float,
min_motion: float,
min_ball_x_frac: float,
ghost_max: int = 18,
hot_motion_frac: float = 0.28,
) -> None:
self.deadzone = deadzone
self.min_motion = min_motion
self.min_ball_x_frac = min_ball_x_frac
self.ghost_max = max(1, int(ghost_max))
self.hot_motion_frac = float(hot_motion_frac)
self._prev_gray: Optional[np.ndarray] = None
self.telemetry: dict[str, Any] = {}
self._track_cx: Optional[float] = None
self._track_cy: Optional[float] = None
self._vx: float = 0.0
self._vy: float = 0.0
self._ghost: int = 0
def reset(self) -> None:
self._prev_gray = None
self.telemetry = {}
self._track_cx = None
self._track_cy = None
self._vx = 0.0
self._vy = 0.0
self._ghost = 0
@staticmethod
def _paddle_py_from_gray(g: np.ndarray, h: int, w: int, y0: int, y1: int) -> Optional[float]:
px0 = int(0.84 * w)
strip = g[y0:y1, px0:]
thr = max(float(np.percentile(strip, 92)), 82.0)
ys = np.where(strip >= thr)[0] + y0
if ys.size == 0:
return None
return float(np.median(ys))
def _effective_min_motion(self) -> float:
if self._track_cx is not None and self._ghost > 0:
return self.min_motion * self.hot_motion_frac
return self.min_motion
def _centroid_from_db(
self, db: np.ndarray, s: float, h: int, w: int, y0: int, y1: int
) -> tuple[float, float]:
Y = np.arange(h, dtype=np.float64)[:, None]
X = np.arange(w, dtype=np.float64)[None, :]
cy = float((db * Y).sum() / s)
cx = float((db * X).sum() / s)
return cx, cy
def _update_track_velocity(self, cx: float, cy: float) -> None:
if self._track_cx is not None and self._track_cy is not None:
self._vx = float(cx - self._track_cx)
self._vy = float(cy - self._track_cy)
self._track_cx, self._track_cy = cx, cy
self._ghost = 0
def _action_from_ball_y(self, ball_y: float, paddle_py: float) -> Optional[int]:
self.telemetry["delta_y"] = float(ball_y - paddle_py)
if ball_y < paddle_py - self.deadzone:
self.telemetry["assist_action"] = 2
return 2
if ball_y > paddle_py + self.deadzone:
self.telemetry["assist_action"] = 3
return 3
self.telemetry["reject"] = "aligned"
return None
def propose(self, rgb: np.ndarray) -> Optional[int]:
g = np.asarray(rgb, dtype=np.uint8).mean(axis=2).astype(np.float32)
h, w = g.shape
y0, y1 = 34, min(194, h - 8)
prev = self._prev_gray
paddle_py = self._paddle_py_from_gray(g, h, w, y0, y1)
self.telemetry = {
"h": h,
"w": w,
"paddle_py": paddle_py,
"ball_cx": None,
"ball_cy": None,
"motion_sum": None,
"assist_action": None,
"reject": None,
"delta_y": None,
"ghost": self._ghost,
"ball_cy_pred": None,
}
self._prev_gray = g.copy()
if prev is None or prev.shape != g.shape:
self.telemetry["reject"] = "no_prev"
return None
d = np.abs(g - prev)
# 패들 열(0.84~) 직전까지 포함: 공이 오른쪽으로 붙으면 예전 bx1=0.82*w 밖으로 차분이 빠지던 문제 완화
bx0, bx1 = int(0.20 * w), min(w - 2, int(0.94 * w))
db = np.zeros_like(d, dtype=np.float64)
db[y0:y1, bx0:bx1] = d[y0:y1, bx0:bx1].astype(np.float64)
s = float(db.sum())
self.telemetry["motion_sum"] = s
eff_min = self._effective_min_motion()
if paddle_py is None:
self.telemetry["reject"] = "no_paddle"
return None
if s >= eff_min:
cx, cy = self._centroid_from_db(db, s, h, w, y0, y1)
self.telemetry["ball_cx"] = cx
self.telemetry["ball_cy"] = cy
if cx < w * self.min_ball_x_frac:
self.telemetry["reject"] = "ball_x"
return None
self._update_track_velocity(cx, cy)
self.telemetry["ghost"] = self._ghost
return self._action_from_ball_y(cy, paddle_py)
# 차분 약함: 직전 궤적으로 짧게 coast
self._ghost = min(self._ghost + 1, self.ghost_max + 5)
self.telemetry["ghost"] = self._ghost
if self._track_cy is None or self._track_cx is None or self._ghost > self.ghost_max:
self.telemetry["reject"] = "low_motion"
return None
if self._track_cx < w * self.min_ball_x_frac * 0.85:
self.telemetry["reject"] = "low_motion"
return None
pred_cy = self._track_cy + self._vy * float(self._ghost)
pred_cy = float(np.clip(pred_cy, float(y0) + 2.0, float(y1) - 2.0))
self.telemetry["ball_cx"] = self._track_cx
self.telemetry["ball_cy"] = self._track_cy
self.telemetry["ball_cy_pred"] = pred_cy
self.telemetry["reject"] = "coast"
return self._action_from_ball_y(pred_cy, paddle_py)
def main() -> None:
ap = argparse.ArgumentParser(description="Pong-v4: local VLM (default) or OpenAI")
ap.add_argument("--backend", choices=("local", "openai"), default="local")
ap.add_argument("--model-path", default=_DEFAULT_LOCAL, help="로컬 HF 모델 디렉터리")
ap.add_argument("--device", default="auto", help="auto|cpu|mps|cuda")
ap.add_argument("--openai-model", default="gpt-4o-mini", help="--backend openai 일 때")
ap.add_argument("--frames", type=int, default=8)
ap.add_argument("--stride", type=int, default=4)
ap.add_argument(
"--vlm-input",
choices=("image", "text"),
default="image",
help="image: 멀티프레임 RGB. text: 추정 볼/패들 수치 히스토리(CSV)만 LLM에 전달.",
)
ap.add_argument("--text-history", type=int, default=24, help="--vlm-input text 일 때 프롬프트에 넣을 최근 행 수")
ap.add_argument("--max-side", type=int, default=160)
ap.add_argument("--jpeg-quality", type=int, default=72)
ap.add_argument("--episodes", type=int, default=3)
ap.add_argument("--dry-run", action="store_true")
ap.add_argument("--no-gui", action="store_true")
ap.add_argument(
"--assist",
choices=("none", "intercept"),
default="none",
help="intercept: 매 step 프레임 차이로 공 y vs 오른쪽 패들 y 정렬(2/3). VLM은 stride마다 갱신되나 수직은 보조가 우선.",
)
ap.add_argument("--assist-deadzone", type=float, default=4.0)
ap.add_argument("--assist-min-motion", type=float, default=100.0)
ap.add_argument("--assist-min-ball-x", type=float, default=0.22, help="공 추정 centroid x 가 화면 너비의 이 비율 미만이면 추적 안 함")
ap.add_argument(
"--assist-ghost-max",
type=int,
default=18,
help="차분이 약해진 뒤 마지막 속도로 맞출 y를 추정하는 최대 step (coast)",
)
ap.add_argument(
"--assist-hot-motion-frac",
type=float,
default=0.28,
help="coast 중 motion_sum 임계값 = min_motion * 이 값 (약한 차분도 다시 잡기)",
)
ap.add_argument(
"--print-track",
action="store_true",
help="공 추정(ball_cx, ball_cy)·패들(paddle_py)·motion_sum·reject·실제 step 액션 등을 주기적으로 stdout 출력",
)
ap.add_argument("--print-track-every", type=int, default=1, help="N step 마다 한 줄 출력 (1=매 step)")
args = ap.parse_args()
local_vlm: Optional[LocalVlmPong] = None
if args.backend == "local" and not args.dry_run:
local_vlm = LocalVlmPong(args.model_path, args.device)
api_key = os.environ.get("OPENAI_API_KEY", "").strip()
base_url = os.environ.get("OPENAI_BASE_URL", "https://api.openai.com/v1").strip()
if args.backend == "openai" and not args.dry_run and not api_key:
print("OPENAI_API_KEY 필요 (--backend openai)", file=sys.stderr)
sys.exit(1)
env = gym.make("Pong-v4", render_mode="rgb_array")
n_act = int(env.action_space.n)
assert n_act == 6
viewer: Optional[AtariQtView] = None
app: Optional[QApplication] = None
if not args.no_gui:
app = QApplication.instance() or QApplication(sys.argv)
viewer = AtariQtView("Pong VLM (local)")
viewer.show()
buf = FrameBuffer(maxlen=max(2, args.frames))
state_hist = StateHistoryBuffer(maxlen=max(2, args.text_history)) if args.vlm_input == "text" else None
last_action = 0
geometry: Optional[InterceptAssist] = None
if args.vlm_input == "text" or args.assist == "intercept" or args.print_track:
geometry = InterceptAssist(
deadzone=args.assist_deadzone,
min_motion=args.assist_min_motion,
min_ball_x_frac=args.assist_min_ball_x,
ghost_max=args.assist_ghost_max,
hot_motion_frac=args.assist_hot_motion_frac,
)
def _fmt(v: Any) -> str:
if v is None:
return "None"
if isinstance(v, float):
return f"{v:.2f}"
return str(v)
for ep in range(args.episodes):
obs, _ = env.reset()
done = False
total = 0.0
last_r = 0.0
buf._buf.clear()
step_in_ep = 0
last_action = 0
if state_hist is not None:
state_hist.clear()
if geometry is not None:
geometry.reset()
while not done:
rgb = env.render()
if rgb is None:
rgb = obs if isinstance(obs, np.ndarray) and obs.ndim == 3 else None
if rgb is not None:
buf.push(rgb)
if viewer is not None:
viewer.show_frame(rgb)
a_int: Optional[int] = None
if geometry is not None and rgb is not None:
a_int = geometry.propose(rgb)
if state_hist is not None:
state_hist.push(step_in_ep, dict(geometry.telemetry), last_r)
use_vlm = (step_in_ep % max(1, args.stride) == 0) and (
(args.vlm_input == "text" and state_hist is not None and state_hist.ready() and rgb is not None)
or (args.vlm_input == "image" and buf.ready())
)
if use_vlm and not args.dry_run:
try:
if args.vlm_input == "text":
assert geometry is not None and state_hist is not None
h0 = int(geometry.telemetry.get("h") or 210)
w0 = int(geometry.telemetry.get("w") or 160)
body = _pong_state_prompt(total, last_r, state_hist.csv_block(), h0, w0)
if args.backend == "local":
assert local_vlm is not None
last_action = local_vlm.predict_text(body)
else:
msgs = _build_openai_messages_text(body)
raw = _call_openai_chat(api_key, base_url, args.openai_model, msgs)
last_action = _parse_action_json(raw)
else:
take = buf.as_list_oldest_first()[-args.frames :]
if args.backend == "local":
assert local_vlm is not None
last_action = local_vlm.predict(take, args.max_side, total, last_r)
else:
b64s = [_resize_jpeg_b64(f, args.max_side, args.jpeg_quality) for f in take]
msgs = _build_openai_messages(b64s, total, last_r)
raw = _call_openai_chat(api_key, base_url, args.openai_model, msgs)
last_action = _parse_action_json(raw)
if not (0 <= last_action < n_act):
last_action = 0
except Exception as e:
print(f"VLM 오류 (NOOP): {e}", file=sys.stderr)
last_action = 0
elif use_vlm and args.dry_run:
last_action = int(np.random.randint(0, n_act))
step_action = last_action
if args.assist == "intercept" and a_int is not None:
step_action = a_int
if args.print_track and geometry is not None and step_in_ep % max(1, args.print_track_every) == 0:
t = geometry.telemetry
print(
"track "
f"ep={ep} step={step_in_ep} "
f"paddle_py={_fmt(t.get('paddle_py'))} "
f"ball_cx={_fmt(t.get('ball_cx'))} ball_cy={_fmt(t.get('ball_cy'))} "
f"motion_sum={_fmt(t.get('motion_sum'))} delta_y={_fmt(t.get('delta_y'))} "
f"ghost={t.get('ghost')} ball_cy_pred={_fmt(t.get('ball_cy_pred'))} "
f"reject={t.get('reject')} assist_a={t.get('assist_action')} "
f"last_vlm={last_action} step_act={step_action}",
flush=True,
)
obs, reward, term, trunc, _ = env.step(step_action)
done = term or trunc
last_r = float(reward)
total += last_r
step_in_ep += 1
print(f"episode {ep + 1}/{args.episodes} return={total:.1f} steps={step_in_ep}")
env.close()
if app is not None:
if viewer is not None:
viewer.close()
app.quit()
if __name__ == "__main__":
main()