wehub-resource-sync
417 lines
17 KiB
Python
417 lines
17 KiB
Python
"""Sudoku scheduler for OpenEnv TextArena Sudoku environment.
|
|
|
|
Reference: TRL openenv_sudoku_grpo.ipynb
|
|
Key features:
|
|
1. Multiple reward functions (empty_cell, valid_move, repetition, progress, correct)
|
|
2. Hints system: parse board, provide guaranteed moves and candidates
|
|
3. Board state tracking with content diff for bounded context
|
|
"""
|
|
import asyncio
|
|
import re
|
|
from collections import defaultdict
|
|
from typing import Any, Dict, List, Optional, Union
|
|
|
|
from swift.rollout.multi_turn import OpenEnvScheduler, multi_turns
|
|
|
|
SUDOKU_SYSTEM_PROMPT = """You are an expert Sudoku player with deep knowledge of logical deduction strategies.
|
|
|
|
## GAME RULES
|
|
1. The puzzle is a 9x9 grid divided into nine 3x3 subgrids (boxes)
|
|
2. Some cells are pre-filled with numbers 1-9
|
|
3. Fill empty cells ('.') with numbers 1-9
|
|
4. Each row, column, and 3x3 box must contain 1-9 without repetition
|
|
5. Cannot overwrite pre-filled cells
|
|
6. Invalid moves result in penalties
|
|
|
|
## HOW TO PLAY
|
|
Output your move in this format: [row col number]
|
|
Example: [3 5 7] means place 7 at row 3, column 5.
|
|
You may reason before your move, but always end with [row col number].
|
|
|
|
## STRATEGIC APPROACH
|
|
- Naked Singles: If a cell has only one possible candidate, fill it immediately.
|
|
- Hidden Singles: If a number can only go in one cell within a row/column/box, place it there.
|
|
- Scanning: Look at each row, column, and box to find where numbers can go.
|
|
|
|
## COMMON PITFALLS
|
|
- Don't guess randomly - Sudoku is pure logic
|
|
- Don't overwrite pre-filled cells
|
|
- Don't repeat a move that was already made
|
|
- Coordinates are 1-indexed (1-9)
|
|
|
|
## BOARD READING
|
|
- Rows labeled R1-R9 (top to bottom)
|
|
- Columns labeled C1-C9 (left to right)
|
|
- Empty cells shown as '.'"""
|
|
|
|
|
|
def _is_valid_board_state(board_str: str) -> bool:
|
|
return 'R1' in board_str and 'R9' in board_str and '|' in board_str
|
|
|
|
|
|
def _parse_board(board_str: str) -> list:
|
|
grid = [[0] * 9 for _ in range(9)]
|
|
if not _is_valid_board_state(board_str):
|
|
return grid
|
|
for line in board_str.split('\n'):
|
|
line_stripped = line.strip()
|
|
if line_stripped and line_stripped[0] == 'R' and len(line_stripped) > 1 and line_stripped[1].isdigit():
|
|
row = int(line_stripped[1]) - 1
|
|
col = 0
|
|
for char in line_stripped[2:]:
|
|
if col >= 9:
|
|
break
|
|
if char == '.':
|
|
grid[row][col] = 0
|
|
col += 1
|
|
elif char.isdigit():
|
|
grid[row][col] = int(char)
|
|
col += 1
|
|
return grid
|
|
|
|
|
|
def _count_filled_cells(board_str: str) -> int:
|
|
grid = _parse_board(board_str)
|
|
return sum(1 for row in grid for cell in row if cell != 0)
|
|
|
|
|
|
def _get_valid_numbers(grid: list, row: int, col: int) -> set:
|
|
if grid[row][col] != 0:
|
|
return set()
|
|
used = set()
|
|
for c in range(9):
|
|
if grid[row][c] != 0:
|
|
used.add(grid[row][c])
|
|
for r in range(9):
|
|
if grid[r][col] != 0:
|
|
used.add(grid[r][col])
|
|
box_row, box_col = 3 * (row // 3), 3 * (col // 3)
|
|
for r in range(box_row, box_row + 3):
|
|
for c in range(box_col, box_col + 3):
|
|
if grid[r][c] != 0:
|
|
used.add(grid[r][c])
|
|
return set(range(1, 10)) - used
|
|
|
|
|
|
def _extract_empty_cells_with_candidates(board_str: str, sort_by_difficulty: bool = True):
|
|
grid = _parse_board(board_str)
|
|
cells = []
|
|
for row in range(9):
|
|
for col in range(9):
|
|
if grid[row][col] == 0:
|
|
candidates = _get_valid_numbers(grid, row, col)
|
|
cells.append((row + 1, col + 1, candidates))
|
|
if sort_by_difficulty:
|
|
cells.sort(key=lambda x: len(x[2]))
|
|
return cells
|
|
|
|
|
|
def _extract_empty_cells(board_str: str) -> list:
|
|
"""Return list of (row, col) tuples for empty cells, 0-indexed."""
|
|
grid = _parse_board(board_str)
|
|
return [(r, c) for r in range(9) for c in range(9) if grid[r][c] == 0]
|
|
|
|
|
|
def _extract_board_only(text: str) -> str:
|
|
if not text:
|
|
return ''
|
|
lines = text.split('\n')
|
|
board_lines = []
|
|
in_board = False
|
|
for line in lines:
|
|
stripped = line.strip()
|
|
if stripped.startswith('C1') or (stripped and stripped[0] == 'R' and len(stripped) > 1
|
|
and stripped[1].isdigit()):
|
|
in_board = True
|
|
if in_board and (stripped.startswith('-') or stripped.startswith('R') or stripped.startswith('C1')):
|
|
board_lines.append(line)
|
|
elif (in_board and stripped and not stripped.startswith('-')
|
|
and not (stripped[0] == 'R' and len(stripped) > 1 and stripped[1].isdigit())):
|
|
break
|
|
return '\n'.join(board_lines) if board_lines else ''
|
|
|
|
|
|
def _make_hints(board_str: str, successful_moves: list, failed_moves: list, difficulty: str = 'easy') -> str:
|
|
parts = []
|
|
all_tried = successful_moves + failed_moves
|
|
if all_tried:
|
|
parts.append(f"\nMOVES ALREADY TRIED (do not repeat): {', '.join(all_tried[:10])}")
|
|
if not board_str or not _is_valid_board_state(board_str):
|
|
return '\n'.join(parts)
|
|
|
|
cells = _extract_empty_cells_with_candidates(board_str, sort_by_difficulty=True)
|
|
if cells:
|
|
guaranteed = []
|
|
other = []
|
|
for r, c, candidates in cells[:10]:
|
|
if len(candidates) == 1:
|
|
guaranteed.append(f'[{r} {c} {list(candidates)[0]}]')
|
|
elif len(candidates) <= 3:
|
|
nums = ','.join(str(n) for n in sorted(candidates))
|
|
other.append(f'({r},{c})->{nums}')
|
|
if guaranteed:
|
|
parts.append(f"\nGUARANTEED MOVES (only one option): {', '.join(guaranteed[:5])}")
|
|
if other:
|
|
parts.append(f"Other options: {' | '.join(other[:5])}")
|
|
|
|
return '\n'.join(parts)
|
|
|
|
|
|
class SudokuScheduler(OpenEnvScheduler):
|
|
"""Sudoku scheduler with multi-reward and hints system.
|
|
|
|
Tracks 5 reward components per trajectory:
|
|
- empty_cell_reward: Did the model target empty cells? (+1/-1)
|
|
- valid_move_reward: Were moves accepted by env? (1.0/-0.5/0.0)
|
|
- repetition_reward: Penalty for repeating moves (exponential)
|
|
- progress_reward: How much of the puzzle was filled (0-1)
|
|
- correct_reward: Environment's reward (0 or 1)
|
|
|
|
Combined reward = sum of all components.
|
|
"""
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
self._last_content_len: Dict[str, int] = {}
|
|
# Per-uuid state tracking
|
|
self._board_states: Dict[str, str] = {}
|
|
self._move_counts: Dict[str, defaultdict] = {}
|
|
self._successful_moves: Dict[str, list] = {}
|
|
self._failed_moves: Dict[str, list] = {}
|
|
self._valid_move_scores: Dict[str, list] = {}
|
|
self._empty_cell_scores: Dict[str, list] = {}
|
|
self._correct_scores: Dict[str, list] = {}
|
|
self._repetition_scores: Dict[str, list] = {}
|
|
self._initial_filled: Dict[str, int] = {}
|
|
self._max_filled: Dict[str, int] = {}
|
|
|
|
async def on_trajectory_start(self, requests):
|
|
"""Initialize env, parse board, compute hints."""
|
|
semaphore = asyncio.Semaphore(getattr(self, 'max_concurrent_envs', 4))
|
|
|
|
async def _init_single(req):
|
|
async with semaphore:
|
|
uuid = req.uuid
|
|
if uuid in self._envs:
|
|
await self._close_and_remove(uuid)
|
|
|
|
row_env_config = (req.data_dict or {}).get('env_config', {}) if hasattr(req, 'data_dict') else {}
|
|
env_config = {**getattr(self, 'env_config_defaults', {}), **row_env_config}
|
|
wrapper = self._create_env(env_config)
|
|
|
|
obs, metadata = await asyncio.to_thread(wrapper.reset)
|
|
system_message = env_config.get('system_message', SUDOKU_SYSTEM_PROMPT)
|
|
|
|
content = self._extract_content(obs)
|
|
self._last_content_len[uuid] = len(content)
|
|
|
|
# Parse initial board state
|
|
board = _extract_board_only(content) if _is_valid_board_state(content) else content
|
|
self._board_states[uuid] = content if _is_valid_board_state(content) else ''
|
|
initial_filled = _count_filled_cells(self._board_states[uuid]) if self._board_states[uuid] else 0
|
|
|
|
# Initialize tracking state
|
|
self._move_counts[uuid] = defaultdict(int)
|
|
self._successful_moves[uuid] = []
|
|
self._failed_moves[uuid] = []
|
|
self._valid_move_scores[uuid] = []
|
|
self._empty_cell_scores[uuid] = []
|
|
self._correct_scores[uuid] = []
|
|
self._repetition_scores[uuid] = []
|
|
self._initial_filled[uuid] = initial_filled
|
|
self._max_filled[uuid] = initial_filled
|
|
|
|
# Build initial message with board + hints
|
|
hints = _make_hints(self._board_states[uuid], [], [])
|
|
user_content = f'{board}{hints}' if board else content
|
|
|
|
from swift.rollout.multi_turn import Messages
|
|
messages = []
|
|
if system_message:
|
|
messages.append({'role': 'system', 'content': system_message})
|
|
messages.append({'role': 'user', 'content': user_content})
|
|
req.messages = messages
|
|
|
|
self._envs[uuid] = wrapper
|
|
self._total_rewards[uuid] = 0.0
|
|
self._step_rewards[uuid] = []
|
|
self._pending_obs[uuid] = None
|
|
|
|
await asyncio.gather(*[_init_single(req) for req in requests])
|
|
|
|
async def _close_and_remove(self, uuid):
|
|
"""Override to clean up all tracking state."""
|
|
await super()._close_and_remove(uuid)
|
|
self._last_content_len.pop(uuid, None)
|
|
self._board_states.pop(uuid, None)
|
|
self._move_counts.pop(uuid, None)
|
|
self._successful_moves.pop(uuid, None)
|
|
self._failed_moves.pop(uuid, None)
|
|
self._valid_move_scores.pop(uuid, None)
|
|
self._empty_cell_scores.pop(uuid, None)
|
|
self._correct_scores.pop(uuid, None)
|
|
self._repetition_scores.pop(uuid, None)
|
|
self._initial_filled.pop(uuid, None)
|
|
self._max_filled.pop(uuid, None)
|
|
|
|
def _extract_content(self, observation: Any) -> str:
|
|
if isinstance(observation, dict):
|
|
messages = observation.get('messages', [])
|
|
if messages:
|
|
return messages[0].get('content', '')
|
|
prompt = observation.get('prompt', '')
|
|
if prompt:
|
|
return prompt
|
|
return str(observation)
|
|
|
|
async def on_turn_end(self, infer_request, response_choice, current_turn):
|
|
"""Parse move, step env, compute multi-reward, generate hints."""
|
|
uuid = infer_request.uuid
|
|
wrapper = self._envs.get(uuid)
|
|
if wrapper is None:
|
|
return {'done': True, 'rollout_infos': {}}
|
|
|
|
action_text = response_choice.message.content
|
|
action_dict = self.parse_action(action_text)
|
|
if action_dict is None:
|
|
# Parse failed: end trajectory with penalty instead of polluting env
|
|
self._total_rewards[uuid] = self._total_rewards.get(uuid, 0.0) - 1.0
|
|
self._step_rewards.setdefault(uuid, []).append(-1.0)
|
|
await self._close_and_remove(uuid)
|
|
return {
|
|
'done': True,
|
|
'rollout_infos': {
|
|
'total_reward': self._total_rewards[uuid],
|
|
'step_rewards': list(self._step_rewards.get(uuid, [])),
|
|
'gym_done': True,
|
|
}
|
|
}
|
|
move = action_dict.get('message', '')
|
|
|
|
# Step environment
|
|
obs, env_reward, done, metadata = await asyncio.to_thread(wrapper.step, action_dict)
|
|
correct_score = float(env_reward or 0.0)
|
|
|
|
# Extract new content (diff from last seen)
|
|
full_content = self._extract_content(obs)
|
|
last_len = self._last_content_len.get(uuid, 0)
|
|
new_content = full_content[last_len:] if len(full_content) > last_len else full_content
|
|
self._last_content_len[uuid] = len(full_content)
|
|
|
|
# Check if env says invalid
|
|
new_content_lower = new_content.lower()
|
|
env_says_invalid = any(kw in new_content_lower
|
|
for kw in ['invalid', 'error', 'cannot', 'already', 'violation', 'lost'])
|
|
|
|
# Check if move targets an empty cell
|
|
if self._board_states.get(uuid):
|
|
empty_cells = _extract_empty_cells(self._board_states[uuid])
|
|
# Convert move coords (1-indexed from model) to 0-indexed for comparison
|
|
move_nums = re.findall(r'\d+', move)
|
|
targets_empty = tuple(int(x) - 1 for x in move_nums[:2]) in empty_cells if len(move_nums) >= 3 else True
|
|
else:
|
|
targets_empty = True
|
|
|
|
# Empty cell reward: +1 if targeted empty, -1 if tried to overwrite
|
|
empty_cell_score = 1.0 if targets_empty else -1.0
|
|
|
|
# Repetition tracking
|
|
is_new_move = self._move_counts[uuid][move] == 0
|
|
repetition_count = self._move_counts[uuid][move]
|
|
self._move_counts[uuid][move] += 1
|
|
repetition_score = -min(2**repetition_count, 10.0) if repetition_count > 0 else 0.0
|
|
|
|
# Valid move score
|
|
is_valid = not env_says_invalid and targets_empty
|
|
if is_valid and is_new_move:
|
|
valid_move_score = 1.0
|
|
self._successful_moves[uuid].append(move)
|
|
elif 'please resubmit' in new_content_lower or 'avoid penalties' in new_content_lower:
|
|
valid_move_score = -0.5
|
|
self._failed_moves[uuid].append(move)
|
|
else:
|
|
valid_move_score = 0.0
|
|
if not is_valid:
|
|
self._failed_moves[uuid].append(move)
|
|
|
|
# Update board state if valid and new content has board
|
|
if is_valid and _is_valid_board_state(new_content):
|
|
self._board_states[uuid] = new_content
|
|
current_filled = _count_filled_cells(new_content)
|
|
if current_filled > self._max_filled[uuid]:
|
|
self._max_filled[uuid] = current_filled
|
|
|
|
# Progress reward
|
|
remaining = 81 - self._initial_filled[uuid]
|
|
if remaining > 0:
|
|
progress_score = (self._max_filled[uuid] - self._initial_filled[uuid]) / remaining
|
|
else:
|
|
progress_score = 1.0
|
|
|
|
# Track all scores
|
|
self._valid_move_scores[uuid].append(valid_move_score)
|
|
self._empty_cell_scores[uuid].append(empty_cell_score)
|
|
self._correct_scores[uuid].append(correct_score)
|
|
self._repetition_scores[uuid].append(repetition_score)
|
|
|
|
combined_reward = (
|
|
sum(self._empty_cell_scores[uuid]) / max(len(self._empty_cell_scores[uuid]), 1)
|
|
+ sum(self._valid_move_scores[uuid]) / max(len(self._valid_move_scores[uuid]), 1)
|
|
+ sum(self._repetition_scores[uuid]) / max(len(self._repetition_scores[uuid]), 1) + progress_score
|
|
+ correct_score)
|
|
|
|
self._total_rewards[uuid] = combined_reward
|
|
self._step_rewards.setdefault(uuid, []).append(combined_reward)
|
|
|
|
# Build next observation with board + hints
|
|
if not done:
|
|
board_str = self._board_states.get(uuid, '')
|
|
board = _extract_board_only(board_str) if board_str else ''
|
|
hints = _make_hints(
|
|
board_str,
|
|
self._successful_moves[uuid],
|
|
self._failed_moves[uuid],
|
|
)
|
|
step_num = len(self._successful_moves[uuid])
|
|
next_obs = f'Step {step_num}. Progress: {step_num} cells filled.\n\nBoard:\n{board}{hints}'
|
|
else:
|
|
next_obs = None
|
|
|
|
self._pending_obs[uuid] = next_obs
|
|
|
|
rollout_infos = {
|
|
'total_reward': self._total_rewards[uuid],
|
|
'step_rewards': list(self._step_rewards.get(uuid, [])),
|
|
'gym_done': done,
|
|
'empty_cell_reward': sum(self._empty_cell_scores[uuid]) / max(len(self._empty_cell_scores[uuid]), 1),
|
|
'valid_move_reward': sum(self._valid_move_scores[uuid]) / max(len(self._valid_move_scores[uuid]), 1),
|
|
'repetition_reward': sum(self._repetition_scores[uuid]) / max(len(self._repetition_scores[uuid]), 1),
|
|
'progress_reward': progress_score,
|
|
'correct_reward': correct_score,
|
|
}
|
|
if done:
|
|
await self._close_and_remove(uuid)
|
|
|
|
return {'done': done, 'rollout_infos': rollout_infos}
|
|
|
|
def parse_action(self, text: str) -> Optional[Dict[str, Any]]:
|
|
"""Extract [row col number] from model output. Returns None if parse fails."""
|
|
match = re.search(r'\[\s*(\d+)\s+(\d+)\s+(\d+)\s*\]', text)
|
|
if match:
|
|
row, col, num = match.groups()
|
|
return {'message': f'[{row} {col} {num}]'}
|
|
|
|
numbers = re.findall(r'\d+', text)
|
|
if len(numbers) >= 3:
|
|
return {'message': f'[{numbers[0]} {numbers[1]} {numbers[2]}]'}
|
|
|
|
return None
|
|
|
|
def format_observation(self, observation: Any) -> Union[str, List[Dict]]:
|
|
return self._extract_content(observation)
|
|
|
|
|
|
# Register scheduler so --external_plugins can load it
|
|
|
|
multi_turns['sudoku_scheduler'] = SudokuScheduler
|