【SSL-RL】自监督强化学习： 好奇心驱动探索 (CDE)算法

"""《 CDE算法的实现示例》
    时间：2024.11
    作者：不去幼儿园
"""
import torch
import torch.nn as nn
import torch.optim as optim

# 定义预测模型
class PredictiveModel(nn.Module):
    def __init__(self, state_dim, action_dim):
        super(PredictiveModel, self).__init__()
        self.fc1 = nn.Linear(state_dim + action_dim, 64)
        self.fc2 = nn.Linear(64, state_dim)
        self.relu = nn.ReLU()

    def forward(self, state, action):
        x = torch.cat([state, action], dim=1)
        x = self.relu(self.fc1(x))
        next_state_pred = self.fc2(x)
        return next_state_pred

# 计算内在奖励
def compute_intrinsic_reward(state, action, next_state, model):
    next_state_pred = model(state, action)
    intrinsic_reward = torch.mean((next_state_pred - next_state) ** 2).item()
    return intrinsic_reward

# 更新预测模型
def update_predictive_model(state, action, next_state, model, optimizer):
    next_state_pred = model(state, action)
    loss = torch.mean((next_state_pred - next_state) ** 2)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    return loss.item()

# 示例用法
state_dim = 16
action_dim = 4
predictive_model = PredictiveModel(state_dim, action_dim)
optimizer = optim.Adam(predictive_model.parameters(), lr=1e-3)

# 假设有批量数据
state = torch.randn(64, state_dim)
action = torch.randn(64, action_dim)
next_state = torch.randn(64, state_dim)

# 计算内在奖励
intrinsic_reward = compute_intrinsic_reward(state, action, next_state, predictive_model)
print(f"Intrinsic Reward: {intrinsic_reward}")

# 更新预测模型
loss = update_predictive_model(state, action, next_state, predictive_model, optimizer)
print(f"Predictive Model Loss: {loss}")

# DQN智能体定义
class DQNAgent:
    def __init__(self, state_dim, action_dim, lr=1e-3):
        self.q_network = nn.Sequential(
            nn.Linear(state_dim, 64),
            nn.ReLU(),
            nn.Linear(64, action_dim)
        )
        self.optimizer = optim.Adam(self.q_network.parameters(), lr=lr)

    def select_action(self, state):
        with torch.no_grad():
            q_values = self.q_network(state)
            action = q_values.argmax().item()
        return action

    def update(self, states, actions, rewards, next_states, dones):
        q_values = self.q_network(states).gather(1, actions.unsqueeze(1)).squeeze()
        with torch.no_grad():
            max_next_q_values = self.q_network(next_states).max(1)[0]
            target_q_values = rewards + (0.99 * max_next_q_values * (1 - dones))
        loss = torch.mean((q_values - target_q_values) ** 2)

        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()

# 训练循环
predictive_model = PredictiveModel(state_dim, action_dim)
dqn_agent = DQNAgent(state_dim=state_dim, action_dim=env.action_space.n)
predictive_optimizer = optim.Adam(predictive_model.parameters(), lr=1e-3)

for episode in range(num_episodes):
    state = env.reset()
    done = False
    episode_reward = 0

    while not done:
        state_tensor = torch.tensor(state, dtype=torch.float32).unsqueeze(0)
        
        # 选择动作
        action = dqn_agent.select_action(state_tensor)
        next_state, reward, done, _ = env.step(action)
        
        # 计算内在奖励
        next_state_tensor = torch.tensor(next_state, dtype=torch.float32).unsqueeze(0)
        intrinsic_reward = compute_intrinsic_reward(state_tensor, torch.tensor([[action]], dtype=torch.float32), next_state_tensor, predictive_model)

        # 计算总奖励
        total_reward = reward + beta * intrinsic_reward
        
        # 更新DQN智能体
        dqn_agent.update(state_tensor, torch.tensor([action]), torch.tensor([total_reward]), next_state_tensor, torch.tensor([done]))

        # 更新预测模型
        update_predictive_model(state_tensor, torch.tensor([[action]], dtype=torch.float32), next_state_tensor, predictive_model, predictive_optimizer)

        state = next_state
        episode_reward += reward
    
    print(f"Episode {episode + 1}: Total Reward = {episode_reward}")