與傳統監督學習不同,直接模仿學習在不同時刻所面臨的數據分布可能不同.試設計一個考慮不同時刻數據分布變化的模仿學習算法

import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset
from sklearn.metrics.pairwise import rbf_kernel
from sklearn.neighbors import KernelDensity
import matplotlib.pyplot as pltclass TimeAwareImitationLearning:def __init__(self, state_dim, action_dim, hidden_dim=64, device='cpu'):"""初始化時間感知的模仿學習算法state_dim: 狀態維度action_dim: 動作維度hidden_dim: 隱藏層維度"""self.state_dim = state_dimself.action_dim = action_dimself.device = device# 策略網絡 - 模仿專家行為self.policy = nn.Sequential(nn.Linear(state_dim + 1, hidden_dim),  # +1 是為了包含時間信息nn.ReLU(),nn.Linear(hidden_dim, hidden_dim),nn.ReLU(),nn.Linear(hidden_dim, action_dim)).to(device)# 判別器網絡 - 區分專家和策略生成的軌跡self.discriminator = nn.Sequential(nn.Linear(state_dim + action_dim + 1, hidden_dim),  # +1 是為了包含時間信息nn.ReLU(),nn.Linear(hidden_dim, hidden_dim),nn.ReLU(),nn.Linear(hidden_dim, 1),nn.Sigmoid()).to(device)# 優化器self.policy_optimizer = optim.Adam(self.policy.parameters(), lr=1e-3)self.discriminator_optimizer = optim.Adam(self.discriminator.parameters(), lr=1e-3)# 記錄訓練過程self.train_losses = []def _compute_time_weights(self, expert_times, current_time, sigma=1.0):"""計算時間權重，距離當前時間越近的樣本權重越大"""time_diffs = np.abs(expert_times - current_time)weights = np.exp(-time_diffs / (2 * sigma**2))return weights / np.sum(weights)def _compute_mmd_loss(self, expert_states, policy_states, times, current_time):"""計算最大均值差異(MMD)損失，衡量分布差異"""# 計算時間權重weights = self._compute_time_weights(times, current_time)# 對專家狀態應用時間權重weighted_expert_states = expert_states * weights.reshape(-1, 1)# 計算MMDexpert_kernel = rbf_kernel(weighted_expert_states, weighted_expert_states)policy_kernel = rbf_kernel(policy_states, policy_states)cross_kernel = rbf_kernel(weighted_expert_states, policy_states)mmd = np.mean(expert_kernel) + np.mean(policy_kernel) - 2 * np.mean(cross_kernel)return mmddef train(self, expert_states, expert_actions, expert_times, epochs=100, batch_size=64):"""訓練時間感知的模仿學習模型expert_states: 專家狀態序列 [num_samples, state_dim]expert_actions: 專家動作序列 [num_samples, action_dim]expert_times: 專家時間戳 [num_samples]"""num_samples = expert_states.shape[0]expert_states_tensor = torch.FloatTensor(expert_states).to(self.device)expert_actions_tensor = torch.FloatTensor(expert_actions).to(self.device)expert_times_tensor = torch.FloatTensor(expert_times).reshape(-1, 1).to(self.device)for epoch in range(epochs):# 當前"時間" - 使用訓練輪次的比例作為時間表示current_time = epoch / epochs# 生成策略動作policy_actions = []for i in range(0, num_samples, batch_size):batch_states = expert_states_tensor[i:i+batch_size]batch_times = torch.full((batch_states.shape[0], 1), current_time).to(self.device)policy_action = self.policy(torch.cat([batch_states, batch_times], dim=1))policy_actions.append(policy_action.detach().cpu().numpy())policy_actions = np.vstack(policy_actions)# 計算MMD損失mmd_loss = self._compute_mmd_loss(expert_states, policy_actions, expert_times, current_time)# 訓練判別器for _ in range(5):  # 判別器訓練多次# 隨機采樣批次indices = np.random.randint(0, num_samples, batch_size)batch_expert_states = expert_states_tensor[indices]batch_expert_actions = expert_actions_tensor[indices]batch_expert_times = expert_times_tensor[indices]# 生成策略動作batch_times = torch.full((batch_size, 1), current_time).to(self.device)batch_policy_actions = self.policy(torch.cat([batch_expert_states, batch_times], dim=1))# 計算判別器損失expert_input = torch.cat([batch_expert_states, batch_expert_actions, batch_expert_times], dim=1)policy_input = torch.cat([batch_expert_states, batch_policy_actions, batch_times], dim=1)expert_output = self.discriminator(expert_input)policy_output = self.discriminator(policy_input)# 判別器損失 (最大化區分能力)d_loss = -torch.mean(torch.log(expert_output + 1e-8) + torch.log(1 - policy_output + 1e-8))self.discriminator_optimizer.zero_grad()d_loss.backward()self.discriminator_optimizer.step()# 訓練策略網絡for _ in range(1):  # 策略網絡訓練較少次數indices = np.random.randint(0, num_samples, batch_size)batch_states = expert_states_tensor[indices]batch_times = torch.full((batch_size, 1), current_time).to(self.device)# 生成策略動作actions = self.policy(torch.cat([batch_states, batch_times], dim=1))# 計算策略損失 (最小化判別器的區分能力)policy_input = torch.cat([batch_states, actions, batch_times], dim=1)policy_output = self.discriminator(policy_input)# 策略損失 + MMD正則化p_loss = -torch.mean(torch.log(policy_output + 1e-8)) + 0.1 * mmd_lossself.policy_optimizer.zero_grad()p_loss.backward()self.policy_optimizer.step()# 記錄損失self.train_losses.append(p_loss.item())if epoch % 100 == 0:print(f"Epoch {epoch}, Loss: {p_loss.item():.4f}, MMD: {mmd_loss:.4f}")def predict(self, state, time):"""根據當前狀態和時間預測動作"""state_tensor = torch.FloatTensor(state).reshape(1, -1).to(self.device)time_tensor = torch.FloatTensor([time]).reshape(1, 1).to(self.device)with torch.no_grad():action = self.policy(torch.cat([state_tensor, time_tensor], dim=1))return action.cpu().numpy()[0]def visualize_training(self):"""可視化訓練過程"""plt.figure(figsize=(10, 6))plt.plot(self.train_losses)plt.title('Training Loss')plt.xlabel('Epoch')plt.ylabel('Loss')plt.grid(True)plt.show()# 示例：生成具有時間分布變化的專家數據
def generate_time_varying_expert_data(num_samples=1000, state_dim=2, time_period=1.0):"""生成隨時間變化的數據分布"""times = np.linspace(0, time_period, num_samples)states = []actions = []for t in times:# 狀態分布隨時間變化mean = np.array([np.sin(2 * np.pi * t), np.cos(2 * np.pi * t)])cov = np.diag([0.1 + 0.1 * np.abs(np.sin(np.pi * t)), 0.1 + 0.1 * np.abs(np.cos(np.pi * t))])state = np.random.multivariate_normal(mean, cov)# 動作是狀態的函數，也隨時間變化action = 2.0 * state * (1.0 + 0.5 * np.sin(2 * np.pi * t))states.append(state)actions.append(action)return np.array(states), np.array(actions), times# 測試算法
def test_time_aware_il():# 生成專家數據state_dim = 2action_dim = 2expert_states, expert_actions, expert_times = generate_time_varying_expert_data(num_samples=2000, state_dim=state_dim, time_period=1.0)# 創建并訓練模型model = TimeAwareImitationLearning(state_dim, action_dim)model.train(expert_states, expert_actions, expert_times, epochs=500)# 可視化訓練過程model.visualize_training()# 測試不同時間點的策略test_times = np.linspace(0, 1, 5)test_states = np.random.randn(len(test_times), state_dim)plt.figure(figsize=(12, 8))for i, t in enumerate(test_times):plt.subplot(2, 3, i+1)# 真實專家行為expert_mask = (expert_times >= t - 0.1) & (expert_times <= t + 0.1)plt.scatter(expert_states[expert_mask, 0], expert_states[expert_mask, 1], c='blue', alpha=0.5, label='Expert')# 模型預測行為pred_actions = np.array([model.predict(s, t) for s in expert_states[expert_mask]])plt.scatter(pred_actions[:, 0], pred_actions[:, 1], c='red', alpha=0.5, label='Policy')plt.title(f'Time = {t:.2f}')plt.xlabel('State 1')plt.ylabel('State 2')plt.legend()plt.tight_layout()plt.show()if __name__ == "__main__":test_time_aware_il()