This commit is contained in:
wehub-resource-sync
+52
@@ -0,0 +1,52 @@
|
||||
# ============================================================
|
||||
# Swift GRPO training with OpenEnv TextArena Sudoku
|
||||
#
|
||||
# Prerequisites:
|
||||
# 1. Start Sudoku server (separate terminal):
|
||||
# TEXTARENA_ENV_ID=Sudoku-v0 MAX_CONCURRENT_ENVS=8 \
|
||||
# python examples/train/grpo/plugin/openenv/start_sudoku_server.py
|
||||
#
|
||||
# 2. This script uses colocate mode:
|
||||
# - vLLM and training share the same GPUs
|
||||
# - No separate rollout server needed
|
||||
#
|
||||
# Environment: TextArena Sudoku (local server, port 8000)
|
||||
# Model: Qwen3.5-4B (enable_thinking=false)
|
||||
# Scheduler: SudokuScheduler (multi-turn, content diff tracking)
|
||||
# Multi-turn: max_turns=20 (20 moves per game)
|
||||
# Rewards: 5-component (empty_cell/valid_move/repetition/progress/correct)
|
||||
# Hints: Board parsing + guaranteed moves + candidates
|
||||
#
|
||||
# ============================================================
|
||||
|
||||
CUDA_VISIBLE_DEVICES=0,1,2,3 \
|
||||
NPROC_PER_NODE=4 \
|
||||
swift rlhf \
|
||||
--rlhf_type grpo \
|
||||
--model Qwen/Qwen3.5-4B \
|
||||
--dataset examples/train/grpo/plugin/openenv/sudoku.jsonl#1000 \
|
||||
--external_plugins examples/train/grpo/plugin/openenv/sudoku_scheduler.py \
|
||||
--enable_thinking false \
|
||||
--torch_dtype bfloat16 \
|
||||
--max_completion_length 256 \
|
||||
--max_length 8192 \
|
||||
--learning_rate 5e-6 \
|
||||
--num_train_epochs 3 \
|
||||
--per_device_train_batch_size 1 \
|
||||
--num_generations 4 \
|
||||
--generation_batch_size 4 \
|
||||
--gradient_accumulation_steps 4 \
|
||||
--temperature 1 \
|
||||
--use_vllm true \
|
||||
--vllm_mode colocate \
|
||||
--vllm_max_model_len 12288 \
|
||||
--vllm_gpu_memory_utilization 0.35 \
|
||||
--gradient_checkpointing true \
|
||||
--use_gym_env true \
|
||||
--multi_turn_scheduler sudoku_scheduler \
|
||||
--max_turns 20 \
|
||||
--save_strategy steps \
|
||||
--save_steps 50 \
|
||||
--logging_steps 1 \
|
||||
--log_completions true \
|
||||
--report_to tensorboard swanlab
|
||||
@@ -0,0 +1,65 @@
|
||||
# ============================================================
|
||||
# Swift GRPO training with OpenEnv TextArena Sudoku (Server Mode)
|
||||
#
|
||||
# Prerequisites:
|
||||
# 1. Start Sudoku server (separate terminal):
|
||||
# TEXTARENA_ENV_ID=Sudoku-v0 MAX_CONCURRENT_ENVS=8 \
|
||||
# python examples/train/grpo/plugin/openenv/start_sudoku_server.py
|
||||
#
|
||||
# 2. Start vLLM rollout server (separate terminal):
|
||||
# CUDA_VISIBLE_DEVICES=0 \
|
||||
# swift rollout \
|
||||
# --model Qwen/Qwen3.5-4B \
|
||||
# --external_plugins examples/train/grpo/plugin/openenv/sudoku_scheduler.py \
|
||||
# --enable_thinking false \
|
||||
# --max_length 8192 \
|
||||
# --vllm_max_model_len 12288 \
|
||||
# --vllm_gpu_memory_utilization 0.9 \
|
||||
# --use_gym_env true \
|
||||
# --multi_turn_scheduler sudoku_scheduler \
|
||||
# --max_turns 20
|
||||
#
|
||||
# 3. This script starts training in server mode:
|
||||
# - vLLM rollout server handles multi-turn + env interaction
|
||||
# - Training process sends generation requests to rollout server
|
||||
# - --multi_turn_scheduler / --max_turns go to BOTH rollout and rlhf
|
||||
#
|
||||
# Environment: TextArena Sudoku (local server, port 8000)
|
||||
# Model: Qwen3.5-4B (enable_thinking=false)
|
||||
# Scheduler: SudokuScheduler (multi-turn, content diff tracking)
|
||||
# Multi-turn: max_turns=20 (20 moves per game)
|
||||
# Rewards: 5-component (empty_cell/valid_move/repetition/progress/correct)
|
||||
# Hints: Board parsing + guaranteed moves + candidates
|
||||
#
|
||||
# ============================================================
|
||||
|
||||
CUDA_VISIBLE_DEVICES=1,2,3 \
|
||||
NPROC_PER_NODE=3 \
|
||||
swift rlhf \
|
||||
--rlhf_type grpo \
|
||||
--model Qwen/Qwen3.5-4B \
|
||||
--dataset examples/train/grpo/plugin/openenv/sudoku.jsonl \
|
||||
--external_plugins examples/train/grpo/plugin/openenv/sudoku_scheduler.py \
|
||||
--enable_thinking false \
|
||||
--torch_dtype bfloat16 \
|
||||
--max_completion_length 256 \
|
||||
--max_length 8192 \
|
||||
--learning_rate 5e-6 \
|
||||
--num_train_epochs 3 \
|
||||
--per_device_train_batch_size 1 \
|
||||
--num_generations 6 \
|
||||
--gradient_accumulation_steps 4 \
|
||||
--temperature 1 \
|
||||
--use_vllm true \
|
||||
--vllm_mode server \
|
||||
--vllm_server_host 127.0.0.1 \
|
||||
--vllm_server_port 8001 \
|
||||
--gradient_checkpointing true \
|
||||
--use_gym_env true \
|
||||
--multi_turn_scheduler sudoku_scheduler \
|
||||
--max_turns 20 \
|
||||
--save_strategy steps \
|
||||
--save_steps 50 \
|
||||
--logging_steps 1 \
|
||||
--log_completions true \
|
||||
--report_to tensorboard swanlab
|
||||
@@ -0,0 +1,47 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Start the TextArena Sudoku server with configurable concurrent sessions.
|
||||
|
||||
The default OpenEnv server only allows 1 concurrent session because
|
||||
TextArenaEnvironment is not marked as SUPPORTS_CONCURRENT_SESSIONS.
|
||||
Since each WebSocket session creates an independent game instance,
|
||||
it is safe to enable concurrent sessions.
|
||||
|
||||
Usage:
|
||||
TEXTARENA_ENV_ID=Sudoku-v0 python start_sudoku_server.py
|
||||
TEXTARENA_ENV_ID=Sudoku-v0 MAX_CONCURRENT_ENVS=8 python start_sudoku_server.py
|
||||
"""
|
||||
import os
|
||||
import uvicorn
|
||||
from openenv.core.env_server.http_server import create_app
|
||||
from textarena_env.server.app import (TextArenaAction, TextArenaObservation, build_textarena_gradio_app,
|
||||
create_textarena_environment)
|
||||
from textarena_env.server.environment import TextArenaEnvironment
|
||||
|
||||
# Read config from environment
|
||||
# Note: TEXTARENA_ENV_ID is read by create_textarena_environment factory,
|
||||
# not by this script directly.
|
||||
max_concurrent_envs = int(os.getenv('MAX_CONCURRENT_ENVS', '8'))
|
||||
host = os.getenv('HOST', '0.0.0.0')
|
||||
port = int(os.getenv('PORT', '8000'))
|
||||
|
||||
# Mark TextArenaEnvironment as supporting concurrent sessions.
|
||||
# Each WebSocket session creates an independent game instance via the factory,
|
||||
# so concurrent sessions are safe.
|
||||
|
||||
TextArenaEnvironment.SUPPORTS_CONCURRENT_SESSIONS = True
|
||||
|
||||
# Build the app with custom max_concurrent_envs
|
||||
|
||||
app = create_app(
|
||||
create_textarena_environment,
|
||||
TextArenaAction,
|
||||
TextArenaObservation,
|
||||
env_name='textarena_env',
|
||||
max_concurrent_envs=max_concurrent_envs,
|
||||
gradio_builder=build_textarena_gradio_app,
|
||||
)
|
||||
|
||||
if __name__ == '__main__':
|
||||
env_id = os.getenv('TEXTARENA_ENV_ID', 'Sudoku-v0')
|
||||
print(f'Starting server: env={env_id}, max_concurrent_envs={max_concurrent_envs}')
|
||||
uvicorn.run(app, host=host, port=port)
|
||||
@@ -0,0 +1,10 @@
|
||||
{"messages": [{"role": "user", "content": "Play"}], "env_config": {"name": "openenv", "base_url": "http://127.0.0.1:8000", "reset_kwargs": {}}}
|
||||
{"messages": [{"role": "user", "content": "Play"}], "env_config": {"name": "openenv", "base_url": "http://127.0.0.1:8000", "reset_kwargs": {}}}
|
||||
{"messages": [{"role": "user", "content": "Play"}], "env_config": {"name": "openenv", "base_url": "http://127.0.0.1:8000", "reset_kwargs": {}}}
|
||||
{"messages": [{"role": "user", "content": "Play"}], "env_config": {"name": "openenv", "base_url": "http://127.0.0.1:8000", "reset_kwargs": {}}}
|
||||
{"messages": [{"role": "user", "content": "Play"}], "env_config": {"name": "openenv", "base_url": "http://127.0.0.1:8000", "reset_kwargs": {}}}
|
||||
{"messages": [{"role": "user", "content": "Play"}], "env_config": {"name": "openenv", "base_url": "http://127.0.0.1:8000", "reset_kwargs": {}}}
|
||||
{"messages": [{"role": "user", "content": "Play"}], "env_config": {"name": "openenv", "base_url": "http://127.0.0.1:8000", "reset_kwargs": {}}}
|
||||
{"messages": [{"role": "user", "content": "Play"}], "env_config": {"name": "openenv", "base_url": "http://127.0.0.1:8000", "reset_kwargs": {}}}
|
||||
{"messages": [{"role": "user", "content": "Play"}], "env_config": {"name": "openenv", "base_url": "http://127.0.0.1:8000", "reset_kwargs": {}}}
|
||||
{"messages": [{"role": "user", "content": "Play"}], "env_config": {"name": "openenv", "base_url": "http://127.0.0.1:8000", "reset_kwargs": {}}}
|
||||
@@ -0,0 +1,416 @@
|
||||
"""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
|
||||
Reference in New Issue
Block a user