TD3 (Twin Delayed Deep Deterministic Policy Gradient)の概要とアルゴリズム及び実装例

import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
import random
import copy

# TD3アルゴリズムの実装
class TD3:
    def __init__(self, state_dim, action_dim, max_action):
        # ニューラルネットワークの初期化
        self.actor = Actor(state_dim, action_dim, max_action)
        self.actor_target = copy.deepcopy(self.actor)
        self.actor_optimizer = optim.Adam(self.actor.parameters())

        self.critic = Critic(state_dim, action_dim)
        self.critic_target = copy.deepcopy(self.critic)
        self.critic_optimizer = optim.Adam(self.critic.parameters())

        self.max_action = max_action

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

    def train(self, replay_buffer, iterations, batch_size=100, discount=0.99, tau=0.005, policy_noise=0.2, noise_clip=0.5, policy_freq=2):
        for it in range(iterations):
            # リプレイバッファからランダムにバッチをサンプリング
            batch_states, batch_next_states, batch_actions, batch_rewards, batch_dones = replay_buffer.sample(batch_size)
            state = torch.FloatTensor(batch_states)
            next_state = torch.FloatTensor(batch_next_states)
            action = torch.FloatTensor(batch_actions)
            reward = torch.FloatTensor(batch_rewards)
            done = torch.FloatTensor(batch_dones)

            # クリティックの更新
            next_action = self.actor_target(next_state)
            noise = torch.normal(0, policy_noise, size=next_action.size())
            noise = torch.clamp(noise, -noise_clip, noise_clip)
            next_action += noise
            next_action = torch.clamp(next_action, -self.max_action, self.max_action)

            target_Q1, target_Q2 = self.critic_target(next_state, next_action)
            target_Q = torch.min(target_Q1, target_Q2)
            target_Q = reward + ((1 - done) * discount * target_Q).detach()

            current_Q1, current_Q2 = self.critic(state, action)

            critic_loss = nn.MSELoss()(current_Q1, target_Q) + nn.MSELoss()(current_Q2, target_Q)

            self.critic_optimizer.zero_grad()
            critic_loss.backward()
            self.critic_optimizer.step()

            # アクターの更新
            if it % policy_freq == 0:
                actor_loss = -self.critic.Q1(state, self.actor(state)).mean()
                self.actor_optimizer.zero_grad()
                actor_loss.backward()
                self.actor_optimizer.step()

                # ソフト更新
                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)

                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)

# アクターネットワークの定義
class Actor(nn.Module):
    def __init__(self, state_dim, action_dim, max_action):
        super(Actor, self).__init__()
        self.layer1 = nn.Linear(state_dim, 400)
        self.layer2 = nn.Linear(400, 300)
        self.layer3 = nn.Linear(300, action_dim)
        self.max_action = max_action

    def forward(self, x):
        x = torch.relu(self.layer1(x))
        x = torch.relu(self.layer2(x))
        x = self.max_action * torch.tanh(self.layer3(x))
        return x

# クリティックネットワークの定義
class Critic(nn.Module):
    def __init__(self, state_dim, action_dim):
        super(Critic, self).__init__()
        self.layer1 = nn.Linear(state_dim + action_dim, 400)
        self.layer2 = nn.Linear(400, 300)
        self.layer3 = nn.Linear(300, 1)

    def forward(self, x, u):
        x = torch.relu(self.layer1(torch.cat([x, u], 1)))
        x = torch.relu(self.layer2(x))
        x = self.layer3(x)
        return x

# リプレイバッファの定義
class ReplayBuffer:
    def __init__(self, max_size=1000000):
        self.storage = []
        self.max_size = max_size
        self.ptr = 0

    def add(self, state, next_state, action, reward, done):
        if len(self.storage) == self.max_size:
            self.storage[int(self.ptr)] = (state, next_state, action, reward, done)
            self.ptr = (self.ptr + 1) % self.max_size
        else:
            self.storage.append((state, next_state, action, reward, done))

    def sample(self, batch_size):
        ind = np.random.randint(0, len(self.storage), size=batch_size)
        states, next_states, actions, rewards, dones = [], [], [], [], []
        for i in ind:
            s, s_, a, r, d = self.storage[i]
            states.append(np.array(s, copy=False))
            next_states.append(np.array(s_, copy=False))
            actions.append(np.array(a, copy=False))
            rewards.append(np.array(r, copy=False))
            dones.append(np.array(d, copy=False))
        return np.array(states), np.array(next_states), np.array(actions), np.array(rewards