DDPG算法轻松解决AI医疗机器人技术难题：连续动作控制 原创

class Actor(nn.Module):
    """
    针对DDPG算法的演员网络。
    """
    def __init__(self, state_dim, action_dim, max_action,use_batch_norm):
        """
        初始化演员的策略网络

        :参数state_dim: 状态空间的维度
        :参数action_dim: 动作空间的维度
        :参数max_action: 动作的最大值
        """
        super(Actor, self).__init__()
        self.bn1 = nn.LayerNorm(HIDDEN_LAYERS_ACTOR) if use_batch_norm else nn.Identity()
        self.bn2 = nn.LayerNorm(HIDDEN_LAYERS_ACTOR) if use_batch_norm else nn.Identity()

        self.l1 = nn.Linear(state_dim, HIDDEN_LAYERS_ACTOR)
        self.l2 = nn.Linear(HIDDEN_LAYERS_ACTOR, HIDDEN_LAYERS_ACTOR)
        self.l3 = nn.Linear(HIDDEN_LAYERS_ACTOR, action_dim)
        self.max_action = max_action

    def forward(self, state):
        """
        通过网络正向传播。

        :参数state: 输入状态
        :返回值: 动作
        """

        a = torch.relu(self.bn1(self.l1(state)))
        a = torch.relu(self.bn2(self.l2(a)))
        return self.max_action * torch.tanh(self.l3(a))

class Critic(nn.Module):
    """
    针对DDPG算法的评论家网络。
    """
    def __init__(self, state_dim, action_dim,use_batch_norm):
        """
        初始化评论家的值网络。

        :参数state_dim: 状态空间的维度
        :参数action_dim: 动作空间的维度
        """
        super(Critic, self).__init__()
        self.bn1 = nn.BatchNorm1d(HIDDEN_LAYERS_CRITIC) if use_batch_norm else nn.Identity()
        self.bn2 = nn.BatchNorm1d(HIDDEN_LAYERS_CRITIC) if use_batch_norm else nn.Identity()
        self.l1 = nn.Linear(state_dim + action_dim, HIDDEN_LAYERS_CRITIC)

        self.l2 = nn.Linear(HIDDEN_LAYERS_CRITIC, HIDDEN_LAYERS_CRITIC)
        self.l3 = nn.Linear(HIDDEN_LAYERS_CRITIC, 1)

    def forward(self, state, action):
        """
        通过网络的正向传播。

        :参数state:输入状态
        :参数action: 输入动作
        :返回值: “状态-动作”对的Q-值
        """
        q = torch.relu(self.bn1(self.l1(torch.cat([state, action], 1))))
        q = torch.relu(self.bn2(self.l2(q)))
        return self.l3(q)

class ReplayBuffer:
    def __init__(self, capacity):
        self.buffer = deque(maxlen=capacity)

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

    def sample(self, batch_size):
        return random.sample(self.buffer, batch_size)

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

"""
节选自​​https://github.com/vitchyr/rlkit/blob/master/rlkit/exploration_strategies/ou_strategy.py​​
"""
class OUNoise(object):
    def __init__(self, action_space, mu=0.0, theta=0.15, max_sigma=0.3, min_sigma=0.3, decay_period=100000):
        self.mu           = mu
        self.theta        = theta
        self.sigma        = max_sigma
        self.max_sigma    = max_sigma
        self.min_sigma    = min_sigma
        self.decay_period = decay_period
        self.action_dim   = action_space.shape[0]
        self.low          = action_space.low
        self.high         = action_space.high
        self.reset()

    def reset(self):
        self.state = np.ones(self.action_dim) * self.mu

    def evolve_state(self):
        x  = self.state
        dx = self.theta * (self.mu - x) + self.sigma * np.random.randn(self.action_dim)
        self.state = x + dx
        return self.state

    def get_action(self, action, t=0): 
        ou_state = self.evolve_state()
        self.sigma = self.max_sigma - (self.max_sigma - self.min_sigma) * min(1.0, t / self.decay_period)
        return np.clip(action + ou_state, self.low, self.high)

class DDPG():
    """
    深度确定性策略梯度（DDPG）代理。
    """
    def __init__(self, state_dim, action_dim, max_action,use_batch_norm):
        """
        初始化DDPG算法代理。

        :参数state_dim: 状态空间的维度
        :参数action_dim: 动作空间的维度
        :参数max_action: 动作的最大值
        """
        # [第0步]
        #初始化演员的策略网络
        self.actor = Actor(state_dim, action_dim, max_action,use_batch_norm)
        # 使用与演员的策略网络相同的权重初始化演员目标网络
        self.actor_target = Actor(state_dim, action_dim, max_action,use_batch_norm)
        self.actor_target.load_state_dict(self.actor.state_dict())
        self.actor_optimizer = optim.Adam(self.actor.parameters(), lr=ACTOR_LR)

        #初始化评论家的值网络
        self.critic = Critic(state_dim, action_dim,use_batch_norm)
        #使用与评论家的值网络相同的权重初始化评论家的目标网络
        self.critic_target = Critic(state_dim, action_dim,use_batch_norm)
        self.critic_target.load_state_dict(self.critic.state_dict())
        self.critic_optimizer = optim.Adam(self.critic.parameters(), lr=CRITIC_LR)

        #初始化重放缓冲区
        self.replay_buffer = ReplayBuffer(BUFFER_SIZE)

def select_action(self, state):
        """
        根据当前状态选择一个动作。

        :参数state：当前状态
        :返回值：选择的动作
        """
        state = torch.FloatTensor(state.reshape(1, -1))
        action = self.actor(state).cpu().data.numpy().flatten()
        return action

def train(self, use_target_network,use_batch_norm):
        """
        训练DDPG代理

        :参数use_target_network: 是否使用目标网络
        :参数use_batch_norm: 是否使用批量归一化
        """
        if len(self.replay_buffer) < BATCH_SIZE:
            return

        # [第4步]. 从重放缓冲区中抽取一批样本
        batch = self.replay_buffer.sample(BATCH_SIZE)
        state, action, reward, next_state, done = map(np.stack, zip(*batch))

        state = torch.FloatTensor(state)
        action = torch.FloatTensor(action)
        next_state = torch.FloatTensor(next_state)
        reward = torch.FloatTensor(reward.reshape(-1, 1))
        done = torch.FloatTensor(done.reshape(-1, 1))

        #评论家网络更新#
        if use_target_network:
            target_Q = self.critic_target(next_state, self.actor_target(next_state))
        else:
            target_Q = self.critic(next_state, self.actor(next_state))

        # [第5步]. 计算目标Q-value (y_i)
        target_Q = reward + (1 - done) * GAMMA * target_Q
        current_Q = self.critic(state, action)
        critic_loss = nn.MSELoss()(current_Q, target_Q.detach())

        # [第6步]. 使用梯度下降来更新评论家网络的权重
        #以最小化损失函数
        self.critic_optimizer.zero_grad()
        critic_loss.backward()
        self.critic_optimizer.step()

        #更新演员网络#
        actor_loss = -self.critic(state, self.actor(state)).mean()

        # [第7步]. 使用梯度下降来更新演员网络的权重
        #以最小化损失函数和最大化Q-value => 选择产生最高累积奖励的动作
        self.actor_optimizer.zero_grad()
        actor_loss.backward()
        self.actor_optimizer.step()

        # [第8步]. 更新目标网络
        if use_target_network:
            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)

def train_ddpg(use_target_network, use_batch_norm, num_episodes=NUM_EPISODES):
    """
    训练DDPG代理

    :参数use_target_network: 是否使用目标网络
    :参数use_batch_norm: 是否使用批量归一化
    :参数num_episodes: 需要训练的回合数
    :返回值: 回合奖励列表
    """
    agent = DDPG(state_dim, action_dim, 1,use_batch_norm)

    episode_rewards = []
    noise = OUNoise(env.action_space)

    for episode in range(num_episodes):
        state= env.reset()
        noise.reset()
        episode_reward = 0
        done = False
        step=0
        while not done:
            action_actor = agent.select_action(state)
            action = noise.get_action(action_actor,step) # Add noise for exploration
            next_state, reward, done,_= env.step(action)
            done = float(done) if isinstance(done, (bool, int)) else float(done[0])
            agent.replay_buffer.push(state, action, reward, next_state, done)

            if len(agent.replay_buffer) > BATCH_SIZE:
                agent.train(use_target_network,use_batch_norm)

            state = next_state
            episode_reward += reward
            step+=1

        episode_rewards.append(episode_reward)

        if (episode + 1) % 10 == 0:
            print(f"Episode {episode + 1}: Reward = {episode_reward}")

    return agent, episode_rewards