Spaces:
Sleeping
Sleeping
| """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) | |
| def _pocketed(system: pt.System, ball_id: str) -> bool: | |
| return system.balls[ball_id].state.s == ptc.pocketed | |
| 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 | |