""" 多进程并行训练环境 """ import numpy as np import torch from typing import List, Tuple, Optional from concurrent.futures import ThreadPoolExecutor import threading from game import Game2048 from trainer import Transition, TrainingStats import time class ParallelGameEnv: """ 并行游戏环境 使用多线程同时运行多个游戏实例 """ def __init__(self, num_envs: int = 4): """ 初始化并行环境 Args: num_envs: 并行游戏数量 """ self.num_envs = num_envs self.envs = [Game2048() for _ in range(num_envs)] self.states = [env.get_state() for env in self.envs] self.scores = [env.accumulated_score for env in self.envs] self.situational_scores = [env.situational_score for env in self.envs] # 游戏结束回调 self.on_game_end = None # 已完成游戏计数 self.games_completed = 0 # 使用线程池 self.executor = ThreadPoolExecutor(max_workers=num_envs) def reset(self, indices: Optional[List[int]] = None) -> Tuple[np.ndarray, np.ndarray]: """ 重置指定环境 Args: indices: 要重置的环境索引，None表示全部重置 Returns: states: (num_envs, 4, 4) 状态数组 score_features: (num_envs, 2) 分数特征数组 """ if indices is None: indices = range(self.num_envs) for i in indices: self.states[i] = self.envs[i].reset() self.scores[i] = self.envs[i].accumulated_score self.situational_scores[i] = self.envs[i].situational_score return self._get_batch_state() def reset_single(self, idx: int) -> Tuple[np.ndarray, np.ndarray]: """重置单个环境""" self.states[idx] = self.envs[idx].reset() self.scores[idx] = self.envs[idx].accumulated_score self.situational_scores[idx] = self.envs[idx].situational_score return self.states[idx], self._get_score_features(idx) def step(self, actions: List[int]) -> List[Transition]: """ 并行执行动作 Args: actions: 每个环境要执行的动作列表 Returns: transitions: 状态转移列表 """ transitions = [] for i, action in enumerate(actions): old_state = self.states[i].copy() old_scores = self._get_score_features(i) old_situational = self.situational_scores[i] # 执行动作 new_state, reward, moved, done = self.envs[i].move(action) # 更新状态 self.states[i] = new_state self.scores[i] = self.envs[i].accumulated_score self.situational_scores[i] = self.envs[i].situational_score # 创建转移记录 transition = Transition( state=old_state, scores=old_scores, action=action, reward=reward, next_state=new_state.copy(), next_scores=self._get_score_features(i), done=done, log_prob=0.0, # 需要在外部填充 value=0.0, # 需要在外部填充 valid_actions=self.envs[i].get_valid_actions() ) transitions.append(transition) # 如果游戏结束，记录统计并重置 if done: # 记录游戏统计 game_stats = { 'score': self.scores[i], 'situational_score': self.situational_scores[i], 'max_tile': self.envs[i].get_max_tile(), 'moves': self.envs[i].moves_count } self.games_completed += 1 # 调用回调 if self.on_game_end: self.on_game_end(game_stats) self.reset_single(i) return transitions def _get_batch_state(self) -> Tuple[np.ndarray, np.ndarray]: """获取批量状态""" states = np.array(self.states, dtype=np.float32) score_features = np.array([ self._get_score_features(i) for i in range(self.num_envs) ], dtype=np.float32) return states, score_features def _get_score_features(self, idx: int) -> np.ndarray: """获取单个环境的分数特征""" max_accumulated = 50000 max_situational = 200 return np.array([ min(self.scores[idx] / max_accumulated, 1.0), min(self.situational_scores[idx] / max_situational, 1.0) ], dtype=np.float32) def get_valid_actions(self) -> np.ndarray: """获取所有环境的有效动作""" return np.array([env.get_valid_actions() for env in self.envs]) def get_game_stats(self) -> List[dict]: """获取所有游戏的统计信息""" return [ { 'score': env.accumulated_score, 'situational_score': env.situational_score, 'max_tile': env.get_max_tile(), 'moves': env.moves_count, 'game_over': env.game_over } for env in self.envs ] def close(self): """关闭环境""" self.executor.shutdown(wait=False) class TrainingWorker: """ 训练工作器 负责收集轨迹数据 """ def __init__( self, model, env: ParallelGameEnv, device: str = "cpu" ): self.model = model.to(device) self.model.eval() self.env = env self.device = device self.stats = TrainingStats() def collect_trajectories( self, num_steps: int = 256, deterministic: bool = False ) -> List[Transition]: """ 收集轨迹数据 Args: num_steps: 每个环境收集的步数 deterministic: 是否确定性选择动作 Returns: transitions: 收集的转移数据 """ all_transitions = [] for _ in range(num_steps): # 获取当前状态 states = np.array(self.env.states, dtype=np.float32) score_features = np.array([ self.env._get_score_features(i) for i in range(self.env.num_envs) ], dtype=np.float32) valid_actions = self.env.get_valid_actions() # 转换为张量 states_t = torch.FloatTensor(states).to(self.device) scores_t = torch.FloatTensor(score_features).to(self.device) valid_t = torch.BoolTensor(valid_actions).to(self.device) # 选择动作 actions = [] log_probs = [] values = [] with torch.no_grad(): for i in range(self.env.num_envs): action, log_prob, value = self.model.get_action( states_t[i:i+1], scores_t[i:i+1], valid_t[i:i+1], deterministic=deterministic ) actions.append(action) log_probs.append(log_prob) values.append(value) # 执行动作 transitions = self.env.step(actions) # 填充log_prob和value for i, t in enumerate(transitions): t.log_prob = log_probs[i] t.value = values[i] all_transitions.append(t) return all_transitions def run_episode( self, deterministic: bool = True, max_steps: int = 10000 ) -> dict: """ 运行一局演示游戏 Args: deterministic: 是否确定性选择 max_steps: 最大步数 Returns: 游戏统计信息 """ # 重置环境 env = Game2048() state = env.reset() total_reward = 0 steps = 0 while not env.game_over and steps < max_steps: # 获取分数特征 max_accumulated = 50000 max_situational = 200 scores = np.array([ min(env.accumulated_score / max_accumulated, 1.0), min(env.situational_score / max_situational, 1.0) ], dtype=np.float32) # 获取有效动作 valid_actions = env.get_valid_actions() # 转换为张量 state_t = torch.FloatTensor(state).unsqueeze(0).to(self.device) scores_t = torch.FloatTensor(scores).unsqueeze(0).to(self.device) valid_t = torch.BoolTensor(valid_actions).unsqueeze(0).to(self.device) # 选择动作 with torch.no_grad(): action, _, _ = self.model.get_action( state_t, scores_t, valid_t, deterministic=deterministic ) # 执行动作 state, reward, moved, done = env.move(action) total_reward += reward steps += 1 return { 'score': env.accumulated_score, 'situational_score': env.situational_score, 'max_tile': env.get_max_tile(), 'steps': steps, 'total_reward': total_reward } class TrainingLoop: """ 完整的训练循环 """ def __init__( self, model, trainer, num_envs: int = 4, device: str = "cpu", steps_per_update: int = 256, save_interval: int = 100, checkpoint_dir: str = "checkpoints" ): self.model = model self.trainer = trainer self.num_envs = num_envs self.device = device self.steps_per_update = steps_per_update self.save_interval = save_interval self.checkpoint_dir = checkpoint_dir # 初始化环境和工作器 self.env = ParallelGameEnv(num_envs=num_envs) self.worker = TrainingWorker(model, self.env, device) # 训练状态 self.training = False self.paused = False self.stats = TrainingStats() # 回调函数 self.on_update_callback = None self.on_game_end_callback = None def train( self, total_games: int = 10000, stop_threshold: int = 100, min_improvement: float = 0.01 ) -> None: """ 训练循环 Args: total_games: 总游戏局数 stop_threshold: 无提升停止阈值（局数） min_improvement: 最小提升比例 """ self.training = True games_since_improvement = 0 best_avg_score = 0 # 设置游戏结束回调 def on_game_end(game_stats): self.stats.record_game( score=game_stats['score'], situational_score=game_stats['situational_score'], max_tile=game_stats['max_tile'], steps=game_stats['moves'] ) if self.on_game_end_callback: self.on_game_end_callback(game_stats) self.env.on_game_end = on_game_end try: while self.training and self.env.games_completed < total_games: if self.paused: time.sleep(0.1) continue # 收集轨迹 transitions = self.worker.collect_trajectories( num_steps=self.steps_per_update // self.num_envs, deterministic=False ) # 更新模型 from trainer import RolloutBuffer buffer = RolloutBuffer(capacity=len(transitions)) buffer.push_batch(transitions) update_stats = self.trainer.update(buffer) if self.on_update_callback: self.on_update_callback(update_stats) # 检查停止条件（每10次更新检查一次） if self.env.games_completed % 10 == 0 and self.env.games_completed > 0: current_avg = self.stats.get_avg_stats(window=100)['avg_score'] if current_avg > best_avg_score * (1 + min_improvement): best_avg_score = current_avg games_since_improvement = 0 else: games_since_improvement = self.env.games_completed - int(best_avg_score / 100 * 100) if best_avg_score > 0 else 0 if games_since_improvement >= stop_threshold: print(f"No improvement for {stop_threshold} updates, stopping.") break except KeyboardInterrupt: print("Training interrupted by user.") finally: self.training = False self.env.close() def stop(self) -> None: """停止训练""" self.training = False def pause(self) -> None: """暂停训练""" self.paused = True def resume(self) -> None: """恢复训练""" self.paused = False def get_stats(self) -> dict: """获取当前统计信息""" return self.stats.get_avg_stats() def save_checkpoint(self, path: str) -> None: """保存模型检查点""" torch.save({ 'model_state_dict': self.model.state_dict(), 'optimizer_state_dict': self.trainer.optimizer.state_dict(), 'stats': self.stats.get_avg_stats(), 'games_played': self.stats.games_played }, path) def load_checkpoint(self, path: str) -> None: """加载模型检查点""" checkpoint = torch.load(path, map_location=self.device) self.model.load_state_dict(checkpoint['model_state_dict']) self.trainer.optimizer.load_state_dict(checkpoint['optimizer_state_dict']) if __name__ == "__main__": from model import Game2048Transformer from trainer import PPOTrainer # 测试并行环境 env = ParallelGameEnv(num_envs=4) env.reset() print("Testing parallel environment...") for i in range(10): actions = [np.random.randint(0, 4) for _ in range(env.num_envs)] transitions = env.step(actions) print(f"Step {i}: collected {len(transitions)} transitions") env.close() print("Parallel environment test passed!")