model/DeepcubeA_module.py

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from pytorch_lightning import LightningModule
from model.DNN import DNN
from model.Cube import TARGET_STATE_ONE_HOT


class RelativeMSELoss(nn.Module):
    def forward(self, pred, target):
        return torch.mean(((pred - target) / (target + 1e-8)) ** 2)

# 判断A中是否有和b相同的张量
def row_allclose_mask(A, b, rtol=1e-4, atol=1e-6):
    # 计算逐元素误差
    diff = torch.abs(A - b)  # (B, D)
    tol = atol + rtol * torch.abs(b)  # (D,), 广播自动扩展到 (B, D)
    
    # 满足误差条件的元素掩码
    mask_elements = diff <= tol  # (B, D), bool
    
    # 判断每行是否所有元素都满足条件
    mask_rows = mask_elements.all(dim=1)  # (B,)
    
    return mask_rows


class DeepcubeA(LightningModule):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.learning_rate = config.learning_rate
        self.weight_decay = config.weight_decay
        self.convergence_threshold = config.convergence_threshold
        self.chunk_size = config.chunk_size
        self.converged_checkpoint_dir = config.converged_checkpoint_dir
        self.compile = config.compile
                
        # 输入维度（54个贴纸，每个有6种可能的颜色，使用one-hot编码）
        self.input_dim = 54 * 6
        
        self.model_theta = DNN(self.input_dim, num_residual_blocks=4)  # 训练模型
        self.model_theta_e = DNN(self.input_dim, num_residual_blocks=4).eval()  # 监督模型

        self.target_state = torch.tensor(TARGET_STATE_ONE_HOT, dtype=torch.float32).reshape(1, -1)

        if self.compile:
            self.model_theta = torch.compile(self.model_theta)
            self.model_theta_e = torch.compile(self.model_theta_e)

        self.K = 1
        
        # 损失函数
        self.criterion = nn.MSELoss()
        
        # 保存超参数
        self.save_hyperparameters(config)

    def transfer_batch_to_tensor(self, batch):
        """
        批量将batch中的数据转移到tensor并移动到正确的设备上
        参数:
            batch: 输入的batch数据
        返回:
            处理后的batch字典，包含tensor格式的数据
        """
        batch_dict = {}
        for key, value in batch.items():
            if isinstance(value, torch.Tensor):
                batch_dict[key] = value.to(self.device)
            else:
                batch_dict[key] = torch.tensor(value, device=self.device)
        return batch_dict
        
    def forward(self, x):
        return self.model_theta(x)
        
    def model_step(self, batch):
        # 从batch中获取状态和邻居
        batch_dict = self.transfer_batch_to_tensor(batch)
        states = batch_dict['state']
        neighbor_states = batch_dict['neighbors']

        B, N, D = neighbor_states.shape

        states = F.one_hot(states.long(), num_classes=6).float().view(B, -1)
        neighbor_states = F.one_hot(neighbor_states.long(), num_classes=6).float().view(B*N, -1)
        
        # 分块预测以避免显存不足
        num_chunks = (B * N + self.chunk_size - 1) // self.chunk_size
        chunked_neighbors = torch.chunk(neighbor_states, num_chunks, dim=0)
        
        with torch.no_grad():
            neighbor_costs = []
            for chunk in chunked_neighbors:
                mask = row_allclose_mask(chunk, self.target_state.to(chunk.device))
                cost = self.model_theta_e(chunk)
                cost[mask] = 0.0
                neighbor_costs.append(cost)
            
            # 聚合结果
            neighbor_costs = torch.cat(neighbor_costs, dim=0)
            neighbor_costs = neighbor_costs.view(B, N)
        
        # 计算min[J_theta_e(A(x_i, a)) + 1]
        min_neighbor_cost = neighbor_costs.abs().min(dim=1)[0] + 1
        
        # 使用model_theta预测当前状态的cost
        current_cost = self.model_theta(states)
        
        # 总是计算损失
        loss = self.criterion(current_cost.squeeze(), min_neighbor_cost)
        return loss, current_cost
        
    def training_step(self, batch, batch_idx):
        # 调用model_step获取损失
        loss, _ = self.model_step(batch)
        
        # 记录指标
        self.log('train_loss', loss, on_step=True, on_epoch=True, prog_bar=True)
        
        return loss
        
    def on_validation_epoch_end(self):
        # 获取验证损失
        val_loss = self.trainer.callback_metrics.get('val_loss')
        
        if val_loss is not None and val_loss < self.convergence_threshold:
            self.log('converged', True)
            
            # 保存模型参数到专门的收敛模型目录
            import os
            os.makedirs(self.converged_checkpoint_dir, exist_ok=True)
            checkpoint_path = os.path.join(self.converged_checkpoint_dir, f"converged_model_K_{self.K}.pth")
            torch.save(self.model_theta.state_dict(), checkpoint_path)
            print(f'模型已保存到 {checkpoint_path}')

            # 如果收敛，更新model_theta_e
            self.model_theta_e.load_state_dict(self.model_theta.state_dict())
            
            # 原文中没有找到上一轮训练的模型下一轮是否要继承参数，这里选择完全继承上一轮的参数，因为从头训练开销太大
            # self.model_theta = DNN(self.input_dim, num_residual_blocks=4)
            # if self.compile:
            #     self.model_theta = torch.compile(self.model_theta)
            
            # 停止训练
            self.trainer.should_stop = True

    def on_train_end(self):
        # 检查训练是否正常结束（非early stopping）
        # 只有当训练不是因为converged而停止时，才执行保存操作
        if not self.trainer.callback_metrics.get('converged', False):
            # 获取最后一个epoch的验证损失
            val_loss = self.trainer.callback_metrics.get('val_loss')
            
            # 保存模型参数到专门的收敛模型目录
            import os
            os.makedirs(self.converged_checkpoint_dir, exist_ok=True)
            checkpoint_path = os.path.join(self.converged_checkpoint_dir, f"final_model_K_{self.K}.pth")
            torch.save(self.model_theta.state_dict(), checkpoint_path)
            print(f'训练结束，模型已保存到 {checkpoint_path}')

            # 更新model_theta_e
            self.model_theta_e.load_state_dict(self.model_theta.state_dict())
        
    def validation_step(self, batch, batch_idx):
        # 计算验证损失
        loss, current_cost = self.model_step(batch)
        self.log('val_loss', loss, on_epoch=True, prog_bar=True)
        self.log('val_cost', current_cost.mean(), on_epoch=True)
        return loss
        
    def configure_optimizers(self):
        optimizer = optim.AdamW(
            self.model_theta.parameters(),
            lr=self.learning_rate,
            weight_decay=self.weight_decay
        )
        
        return {'optimizer': optimizer}

    def load_state_dict_theta_e(self, checkpoint_path):
        state_dict = torch.load(checkpoint_path)
        self.model_theta_e.load_state_dict(state_dict)
        self.model_theta_e.zero_output = False