"""PositionPlayEnv — curriculum stage 2a cho PoolCoach. Task: cue + bi 1 + bi 2 (vị trí random, cách nhau >= 4R), vẫn 1 cú đánh. Pot bi 1 sao cho bi cái dừng ở vị trí có CÚ THUẬN đánh tiếp bi 2. Thiết kế chốt 13/07/2026 (Documents/PoolCoach_Stage2_PositionPlay_Design.md): GIỮ NGUYÊN toàn bộ chuỗi reward stage 1 đã phá collapse (pot +1.0, scratch -0.3, gated contact bonus, dense aim ghost-ball, shaping khoảng cách) — chỉ THÊM đúng 1 term position, GATED theo (pot && !scratch): reward += POS_COEF * Q(cue_final, b2_final) # POS_COEF = 0.5 Bài học gated contact bonus (11/07): không gate là agent farm position mà bỏ pot. Pot suông +1.0, pot + shape đẹp +1.5 → thứ tự ưu tiên không đổi. Position quality Q ∈ [0, 1]: với mỗi lỗ p: u = unit(b2 − c); v = unit(p − b2); cos_cut = dot(u, v) Q_angle = max(0, max_p cos_cut) # lỗ thuận nhất; không lỗ nào dương → 0 Q = Q_angle * 1 / (1 + d(c, b2) / POS_DIST_REF) Chọn lỗ theo cos_cut lớn nhất — cùng logic _ghost_aim_point (KHÔNG lỗ gần nhất). dist_factor phạt nhẹ cú xa (d=0.5m → 0.67; d=1m → 0.5; d=2m → 0.33). 2 bi → không có snooker, bỏ qua check obstruction. Obs: Box[0,1]^6 = [cue_x, cue_y, b1_x, b1_y, b2_x, b2_y] / (w, l), clip [0,1]. Action: giữ nguyên Box[-1,1]^4 (phi / V0 / side / vert). Stage 2 mới là lúc spin có việc → log |side|, |vert| vào info (figure spin-usage cho luận văn). Edge cases (design doc §5): - Bi 1 không vào → gate đóng, không position term (như stage 1) - Scratch (kể cả pot) → scratch_penalty, không position term - Bi 2 bị đụng văng → hợp lệ, Q tính trên vị trí bi 2 SAU simulate - Bi 2 vô tình rớt lỗ → Q = 1.0 (combo may mắn), log b2_potted theo dõi - simulation_failed → như stage 1, info đủ key (pos_q=0, b2_potted=0) """ from __future__ import annotations import numpy as np import gymnasium as gym from gymnasium import spaces import pooltool as pt import pooltool.constants as ptc BALL_R = 0.028575 # bán kính bi (m), mặc định pooltool # --- Reward config: chuỗi stage 1 giữ NGUYÊN (không đụng vào) --- POT_REWARD = 1.0 # bi 1 vào lỗ SCRATCH_PENALTY = -0.3 # bi cái vào lỗ — kể cả khi bi 1 cũng vào CONTACT_BONUS = 0.2 # gated theo năng lượng truyền (chống tap farming) CONTACT_DISP_REF = 0.5 # (m) quãng đường bi 1 để ăn TRỌN contact bonus AIM_COEF = 0.1 # dense aim reward về ghost-ball của BI 1 (không đổi) # --- Term MỚI duy nhất của stage 2a --- POS_COEF = 0.5 # hệ số position term, GATED theo (pot && !scratch) # Smoke test hỏi đúng 2 câu: pot% giữ ~15-20%? Q có tăng? # Pot sập → giảm 0.5 → 0.25 (position term lấn multi-objective). POS_DIST_REF = 1.0 # (m) tham chiếu dist_factor: d=1m → factor 0.5 def table_bounds(w, l, pockets): """Hộp bao của (mặt bàn ∪ mọi đường tròn lỗ) — trả (lo_x, hi_x, lo_y, hi_y). ``pockets``: iterable object có ``.center`` (≥2 chiều) và ``.radius``, tức đúng shape của ``pt.Table.pockets.values()``. Đây là ENVELOPE VẬT LÝ của mọi vị trí nghỉ hợp lệ: bi CHƯA bị ăn không thể nằm ngoài hộp này, vì muốn ra xa hơn nó phải đi xuyên qua một đường tròn lỗ, mà đi qua đường tròn lỗ là bị ăn. Tâm lỗ nằm NGOÀI hộp bàn (lỗ góc lệch ``depth/√2`` theo đường chéo, lỗ giữa lệch ``depth`` theo phương ngang) nên hộp này rộng hơn ``[0, w] × [0, l]`` thật sự, không phải chỉ trên lý thuyết. Tách thành hàm THUẦN (không đụng pooltool, không đụng self) có chủ đích: envelope phải KIỂM ĐƯỢC là dẫn ra từ hình học chứ không phải hằng số gõ tay — mà `tests/conftest.py` cố ý nhét stub pooltool rỗng nên test không dựng nổi bàn thật. Hàm thuần thì test nạp bàn giả vào được. """ return ( min(0.0, min(float(p.center[0]) - float(p.radius) for p in pockets)), max(float(w), max(float(p.center[0]) + float(p.radius) for p in pockets)), min(0.0, min(float(p.center[1]) - float(p.radius) for p in pockets)), max(float(l), max(float(p.center[1]) + float(p.radius) for p in pockets)), ) class PositionPlayEnv(gym.Env): """Một cú đánh: pot bi 1 + điều bi cái về vị trí thuận cho bi 2.""" metadata = {"render_modes": []} def __init__( self, seed: int | None = None, scratch_penalty: float = SCRATCH_PENALTY, pos_coef: float = POS_COEF, aim_mode: str = "best_cut", ): """pos_coef: override cho ablation POS_COEF ∈ {0, 0.5, 1.0}; pos_coef=0 chính là stage-1-trên-env-mới (position-blind baseline). aim_mode: "best_cut" (mặc định, = stage 1): aim reward về ghost-ball của MỘT lỗ có cut angle thuận nhất — phi bị ghim vào 1 đường pot. "any": aim reward = max cos trên ghost-ball của MỌI lỗ khả thi (dot(cue→bi, bi→lỗ) > 0) — agent tự do chọn lỗ, mở khoá phi cho position play (thí nghiệm 14/07: aim best_cut nghi là nút thắt khiến Q phẳng — phi khoá thì chỉ còn V0/spin tự do). """ super().__init__() assert aim_mode in ("best_cut", "any"), aim_mode self.scratch_penalty = float(scratch_penalty) self.pos_coef = float(pos_coef) self.aim_mode = aim_mode self.table = pt.Table.default() self.w, self.l = self.table.w, self.table.l self._diag = float(np.hypot(self.w, self.l)) self._pockets = [ np.asarray(p.center[:2], dtype=np.float64) for p in self.table.pockets.values() ] # Envelope vị trí nghỉ hợp lệ, tính MỘT LẦN lúc dựng bàn — `validate_full` # chỉ đọc thuộc tính này, không tính lại mỗi cú. Xem `table_bounds`. self._bounds = table_bounds(self.w, self.l, self.table.pockets.values()) self.observation_space = spaces.Box(0.0, 1.0, shape=(6,), dtype=np.float32) self.action_space = spaces.Box(-1.0, 1.0, shape=(4,), dtype=np.float32) self._rng = np.random.default_rng(seed) self._system: pt.System | None = None # ------------------------------------------------------------------ utils def _random_xy(self) -> np.ndarray: margin = 4 * BALL_R return np.array( [ self._rng.uniform(margin, self.w - margin), self._rng.uniform(margin, self.l - margin), ] ) def _min_pocket_dist(self, xy: np.ndarray) -> float: return min(float(np.linalg.norm(xy - p)) for p in self._pockets) def _ghost_aim_point(self, cue_xy: np.ndarray, tgt_xy: np.ndarray) -> np.ndarray: """Điểm ghost-ball cho BI 1 — copy nguyên từ stage 1 (không đổi).""" to_tgt = tgt_xy - cue_xy to_tgt = to_tgt / np.linalg.norm(to_tgt) best_dir, best_score = None, -np.inf for p in self._pockets: to_pocket = p - tgt_xy norm = float(np.linalg.norm(to_pocket)) if norm < 1e-9: continue to_pocket = to_pocket / norm score = float(np.dot(to_tgt, to_pocket)) if score > best_score: best_dir, best_score = to_pocket, score return tgt_xy - 2.0 * BALL_R * best_dir def _ghost_dirs_any(self, cue_xy: np.ndarray, tgt_xy: np.ndarray) -> list: """Hướng cue→ghost cho MỌI lỗ khả thi (aim_mode="any"). Lỗ khả thi: dot(cue→bi, bi→lỗ) > 0 — ghost nằm phía bi cái với tới được. Lỗ ngược hướng bị loại (ghost sau lưng bi, cú bất khả thi; thưởng aim vào đó là thưởng cú trượt). Không lỗ nào khả thi (hiếm) → fallback về best_cut. """ to_tgt = tgt_xy - cue_xy to_tgt = to_tgt / np.linalg.norm(to_tgt) dirs = [] for p in self._pockets: v = p - tgt_xy norm = float(np.linalg.norm(v)) if norm < 1e-9: continue v = v / norm if float(np.dot(to_tgt, v)) <= 0.0: continue g = (tgt_xy - 2.0 * BALL_R * v) - cue_xy g_norm = float(np.linalg.norm(g)) if g_norm < 1e-9: continue dirs.append(g / g_norm) if not dirs: g = self._ghost_aim_point(cue_xy, tgt_xy) - cue_xy dirs.append(g / np.linalg.norm(g)) return dirs def _position_q(self, cue_xy: np.ndarray, b2_xy: np.ndarray) -> float: """Q ∈ [0,1]: chất lượng vị trí bi cái cho cú tiếp theo vào bi 2. Q_angle = max(0, max_p cos_cut) — lỗ thuận nhất theo cut angle, cùng triết lý chọn lỗ với _ghost_aim_point. dist_factor = 1/(1 + d/POS_DIST_REF) — phạt nhẹ cú xa. """ d = float(np.linalg.norm(b2_xy - cue_xy)) if d < 1e-9: # trùng vị trí (không xảy ra thực tế) return 0.0 u = (b2_xy - cue_xy) / d # hướng cú đánh bi 2 best_cos = -np.inf for p in self._pockets: to_pocket = p - b2_xy norm = float(np.linalg.norm(to_pocket)) if norm < 1e-9: continue best_cos = max(best_cos, float(np.dot(u, to_pocket / norm))) q_angle = max(0.0, best_cos) dist_factor = 1.0 / (1.0 + d / POS_DIST_REF) return q_angle * dist_factor def _ball_xy(self, ball_id: str) -> np.ndarray: return np.asarray( self._system.balls[ball_id].state.rvw[0][:2], dtype=np.float64 ) def _obs(self) -> np.ndarray: cue, b1, b2 = self._ball_xy("cue"), self._ball_xy("1"), self._ball_xy("2") obs = np.array( [ cue[0] / self.w, cue[1] / self.l, b1[0] / self.w, b1[1] / self.l, b2[0] / self.w, b2[1] / self.l, ], dtype=np.float32, ) # bi đã vào lỗ nằm ở tâm lỗ (hơi ngoài mép bàn) -> clip về [0, 1] return np.clip(obs, 0.0, 1.0) @staticmethod def _pocketed(system: pt.System, ball_id: str) -> bool: return system.balls[ball_id].state.s == ptc.pocketed @staticmethod def _cue_hit_target(system: pt.System) -> bool: """True nếu có va chạm bi-bi giữa 'cue' và '1' trong lượt mô phỏng.""" for ev in pt.events.filter_type(system.events, pt.EventType.BALL_BALL): if "cue" in ev.ids and "1" in ev.ids: return True return False # ------------------------------------------------------------------ gym API def reset(self, *, seed: int | None = None, options=None): super().reset(seed=seed) if seed is not None: self._rng = np.random.default_rng(seed) # 3 bi, đôi một cách nhau > 4R placed = [self._random_xy()] while len(placed) < 3: xy = self._random_xy() if all(np.linalg.norm(xy - q) > 4 * BALL_R for q in placed): placed.append(xy) cue_xy, b1_xy, b2_xy = placed balls = { "cue": pt.Ball.create("cue", xy=tuple(cue_xy)), "1": pt.Ball.create("1", xy=tuple(b1_xy)), "2": pt.Ball.create("2", xy=tuple(b2_xy)), } self._system = pt.System( table=self.table, balls=balls, cue=pt.Cue(cue_ball_id="cue") ) return self._obs(), {} def step(self, action): action = np.clip(np.asarray(action, dtype=np.float64), -1.0, 1.0) phi = float((action[0] + 1.0) / 2.0 * 360.0) % 360.0 v0 = float(0.5 + (action[1] + 1.0) / 2.0 * 3.5) side = float(action[2] * 0.4) vert = float(action[3] * 0.4) abs_side, abs_vert = abs(side) / 0.4, abs(vert) / 0.4 # chuẩn hoá [0,1] # --- dense aim reward: về ghost-ball BI 1, tính TRƯỚC simulate --- cue_xy, tgt_xy = self._ball_xy("cue"), self._ball_xy("1") phi_rad = np.radians(phi) aim_dir = np.array([np.cos(phi_rad), np.sin(phi_rad)]) if self.aim_mode == "any": # max trên mọi lỗ khả thi — agent tự chọn lỗ, phi không bị ghim aim_cos = max( float(np.dot(aim_dir, d)) for d in self._ghost_dirs_any(cue_xy, tgt_xy) ) else: # "best_cut" — nguyên bản stage 1 ghost = self._ghost_aim_point(cue_xy, tgt_xy) to_ghost = ghost - cue_xy to_ghost /= np.linalg.norm(to_ghost) # |cue-ghost| >= 2R > 0 aim_cos = float(np.dot(aim_dir, to_ghost)) aim_r = AIM_COEF * aim_cos d_before = self._min_pocket_dist(tgt_xy) self._system.cue.set_state(V0=v0, phi=phi, a=side, b=vert) try: pt.simulate(self._system, inplace=True) except Exception: # hiếm: edge case vật lý -> coi như foul return self._obs(), self.scratch_penalty + aim_r, True, False, { "potted": 0, "scratch": 0, "contact": 0, "aim_cos": aim_cos, "tgt_disp": 0.0, "pos_q": 0.0, "b2_potted": 0, "abs_side": abs_side, "abs_vert": abs_vert, "error": "simulation_failed", } scratch = self._pocketed(self._system, "cue") potted = self._pocketed(self._system, "1") b2_potted = self._pocketed(self._system, "2") contact = self._cue_hit_target(self._system) tgt_disp = float(np.linalg.norm(self._ball_xy("1") - tgt_xy)) # --- position quality: chỉ tính khi gate mở (pot && !scratch) --- pos_q = 0.0 if potted and not scratch: if b2_potted: pos_q = 1.0 # combo may mắn — theo dõi tần suất qua b2_potted else: pos_q = self._position_q(self._ball_xy("cue"), self._ball_xy("2")) if scratch: reward = self.scratch_penalty elif potted: reward = POT_REWARD + self.pos_coef * pos_q else: d_after = self._min_pocket_dist(self._ball_xy("1")) reward = (d_before - d_after) / self._diag if contact: reward += CONTACT_BONUS * min(1.0, tgt_disp / CONTACT_DISP_REF) reward += aim_r # áp lên MỌI cú — no-op không né được # ép int/float để VecMonitor(info_keywords=...) ghi vào monitor.csv info = { "potted": int(potted), "scratch": int(scratch), "contact": int(contact), "aim_cos": aim_cos, "tgt_disp": tgt_disp, "pos_q": float(pos_q), "b2_potted": int(b2_potted), "abs_side": abs_side, "abs_vert": abs_vert, } return self._obs(), float(reward), True, False, info