用pytorch写添加icm的ddpg算法并添加中文注释

首先，我们需要导入所需要的库。以下是代码的第一部分：

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import gym
from collections import deque
import random
import math

接下来，我们需要定义我们的Actor和Critic神经网络。Actor网络的作用是根据当前状态输出动作，而Critic网络的作用是根据当前状态和动作输出一个Q值。以下是代码的第二部分：

class Actor(nn.Module):
    def __init__(self, state_dim, action_dim, max_action):
        super(Actor, self).__init__()
        self.layer_1 = nn.Linear(state_dim, 256)
        self.layer_2 = nn.Linear(256, 256)
        self.layer_3 = nn.Linear(256, action_dim)
        self.max_action = max_action

    def forward(self, x):
        x = F.relu(self.layer_1(x))
        x = F.relu(self.layer_2(x))
        x = self.max_action * torch.tanh(self.layer_3(x))
        return x


class Critic(nn.Module):
    def __init__(self, state_dim, action_dim):
        super(Critic, self).__init__()

        # Q1 architecture
        self.l1 = nn.Linear(state_dim + action_dim, 256)
        self.l2 = nn.Linear(256, 256)
        self.l3 = nn.Linear(256, 1)

        # Q2 architecture
        self.l4 = nn.Linear(state_dim + action_dim, 256)
        self.l5 = nn.Linear(256, 256)
        self.l6 = nn.Linear(256, 1)

    def forward(self, state, action):
        sa = torch.cat([state, action], 1)

        q1 = F.relu(self.l1(sa))
        q1 = F.relu(self.l2(q1))
        q1 = self.l3(q1)

        q2 = F.relu(self.l4(sa))
        q2 = F.relu(self.l5(q2))
        q2 = self.l6(q2)
        return q1, q2

接下来，我们需要定义我们的ICM模型。ICM模型由一个前向模型和一个反向模型组成，前向模型的作用是预测下一个状态，而反向模型的作用是预测动作。以下是代码的第三部分：

class ICM(nn.Module):
    def __init__(self, state_dim, action_dim):
        super(ICM, self).__init__()
        self.encoder = nn.Sequential(
            nn.Linear(state_dim + action_dim, 256),
            nn.ReLU(),
            nn.Linear(256, 256),
            nn.ReLU(),
            nn.Linear(256, 256),
            nn.ReLU(),
            nn.Linear(256, state_dim)
        )

        self.forward_model = nn.Sequential(
            nn.Linear(state_dim + action_dim, 256),
            nn.ReLU(),
            nn.Linear(256, 256),
            nn.ReLU(),
            nn.Linear(256, 256),
            nn.ReLU(),
            nn.Linear(256, state_dim)
        )

        self.inverse_model = nn.Sequential(
            nn.Linear(state_dim * 2, 256),
            nn.ReLU(),
            nn.Linear(256, 256),
            nn.ReLU(),
            nn.Linear(256, 256),
            nn.ReLU(),
            nn.Linear(256, action_dim)
        )

    def forward(self, state, action, next_state):
        encoded_state_action = torch.cat([state, action], 1)
        encoded_next_state = self.encoder(encoded_state_action)
        predicted_next_state = self.forward_model(encoded_state_action)
        predicted_action = self.inverse_model(torch.cat([state, next_state], 1))
        return encoded_next_state, predicted_next_state, predicted_action

接下来，我们需要定义我们的DDPG算法。DDPG算法由Actor网络、Critic网络和ICM模型组成。以下是代码的第四部分：

class DDPG(object):
    def __init__(self, state_dim, action_dim, max_action):
        self.actor = Actor(state_dim, action_dim, max_action).to(device)
        self.actor_target = Actor(state_dim, action_dim, max_action).to(device)
        self.actor_optimizer = torch.optim.Adam(self.actor.parameters(), lr=1e-3)

        self.critic = Critic(state_dim, action_dim).to(device)
        self.critic_target = Critic(state_dim, action_dim).to(device)
        self.critic_optimizer = torch.optim.Adam(self.critic.parameters(), lr=1e-3)

        self.icm = ICM(state_dim, action_dim).to(device)
        self.icm_optimizer = torch.optim.Adam(self.icm.parameters(), lr=1e-4)

        self.max_action = max_action

    def select_action(self, state):
        state = torch.FloatTensor(state.reshape(1, -1)).to(device)
        return self.actor(state).cpu().data.numpy().flatten()

    def train(self, replay_buffer, iterations, batch_size=100, discount=0.99, tau=0.005, actor_noise=0.1):
        for it in range(iterations):
            # Sample replay buffer
            x, y, u, r, d = replay_buffer.sample(batch_size)
            state = torch.FloatTensor(x).to(device)
            action = torch.FloatTensor(u).to(device)
            next_state = torch.FloatTensor(y).to(device)
            done = torch.FloatTensor(1 - d).to(device)
            reward = torch.FloatTensor(r).to(device)

            # Calculate intrinsic reward
            encoded_next_state, predicted_next_state, predicted_action = self.icm(state, action, next_state)
            intrinsic_reward = ((encoded_next_state - predicted_next_state) ** 2).sum(1) + \
                               ((predicted_action - action) ** 2).sum(1)
            intrinsic_reward = intrinsic_reward.unsqueeze(1)

            # Calculate Q values
            target_Q1, target_Q2 = self.critic_target(next_state, self.actor_target(next_state))
            target_Q = torch.min(target_Q1, target_Q2)
            target_Q = reward + (done * discount * target_Q) + intrinsic_reward

            # Calculate critic loss
            current_Q1, current_Q2 = self.critic(state, action)
            critic_loss = F.mse_loss(current_Q1, target_Q) + F.mse_loss(current_Q2, target_Q)

            # Optimize the critic
            self.critic_optimizer.zero_grad()
            critic_loss.backward()
            self.critic_optimizer.step()

            # Calculate actor loss
            actor_loss = -self.critic(state, self.actor(state))[0].mean()

            # Optimize the actor
            self.actor_optimizer.zero_grad()
            actor_loss.backward()
            self.actor_optimizer.step()

            # Update the frozen target models
            for param, target_param in zip(self.critic.parameters(), self.critic_target.parameters()):
                target_param.data.copy_(tau * param.data + (1 - tau) * target_param.data)

            for param, target_param in zip(self.actor.parameters(), self.actor_target.parameters()):
                target_param.data.copy_(tau * param.data + (1 - tau) * target_param.data)

            # Optimize the ICM
            encoded_next_state_pred, predicted_next_state_pred, predicted_action_pred = self.icm(state, action, next_state)
            forward_loss = F.mse_loss(encoded_next_state_pred, predicted_next_state)
            inverse_loss = F.mse_loss(predicted_action_pred, action)
            icm_loss = forward_loss + inverse_loss

            self.icm_optimizer.zero_grad()
            icm_loss.backward()
            self.icm_optimizer.step()

最后，我们需要定义我们的经验回放缓冲区。经验回放缓冲区的作用是存储并随机采样之前的经验，以便我们可以训练我们的DDPG算法。以下是代码的第五部分：

class ReplayBuffer(object):
    def __init__(self, max_size=1000000):
        self.buffer = deque(maxlen=max_size)

    def add(self, state, next_state, action, reward, done):
        self.buffer.append((state, next_state, action, reward, done))

    def sample(self, batch_size):
        state, next_state, action, reward, done = zip(*random.sample(self.buffer, batch_size))
        return np.array(state), np.array(next_state), np.array(action), np.array(reward).reshape(-1, 1), np.array(done).reshape(-1, 1)

    def __len__(self):
        return len(self.buffer)

现在，我们已经成功地编写了一个添加ICM的DDPG算法，并添加了中文注释。我们可以使用此算法来解决各种强化学习问题。