DRL在计算机视觉、机器学习等领域的应用 Deep Reinforcement Learning for Atari Games

    import torch
    import torch.nn as nn
    import gym
    from collections import deque
    import numpy as np
    import matplotlib.pyplot as plt
    %matplotlib inline
    class DQN(nn.Module):
    def __init__(self, num_inputs, num_outputs):
      super(DQN, self).__init__()
       self.fc = nn.Sequential(
      nn.Linear(num_inputs, 128),
      nn.ReLU(),
      nn.Linear(128, 128),
      nn.ReLU(),
      nn.Linear(128, num_outputs)
      )
    def forward(self, x):
      return self.fc(x)
    env = gym.make('CartPole-v0')
    # set up the model
    model = DQN(env.observation_space.shape[0], env.action_space.n)
    optimizer = torch.optim.Adam(model.parameters())
    loss_fn = nn.MSELoss()
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    print("Using device:", device)
    model = model.to(device)
    def train(model, optimizer, loss_fn, experience):
    state, action, next_state, reward, done = experience
    state = torch.FloatTensor(np.float32(state)).unsqueeze(0).to(device)
    next_state = torch.FloatTensor(np.float32(next_state)).unsqueeze(0).to(device)
    action = torch.LongTensor(action).unsqueeze(0).to(device)
    reward = torch.FloatTensor([reward]).unsqueeze(0).to(device)
    done = torch.FloatTensor([done]).unsqueeze(0).to(device)
    # get Q(s',a) and best action a' using target net
    q_values_next = model(next_state)
    _, actions_next = q_values_next.max(dim=1)
    q_values_next_target = target_net(next_state)
    q_value_next_target = q_values_next_target.gather(1, actions_next.unsqueeze(1)).squeeze(-1)
    # compute target value y
    target = (q_value_next * GAMMA) + (reward + (GAMMA ** N_STEP) * q_value_next_target * (not done))
    # predict current Q values
    q_values = model(state)
    q_value = q_values.gather(1, action.unsqueeze(1)).squeeze(-1)
    # calculate loss between predicted Q value and actual label
    loss = loss_fn(q_value, target)
    # optimize the model
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    def run_episode():
    episode_rewards = []
    state = env.reset()
    while True:
      # select an action based on epsilon greedy strategy
      eps_threshold = EPS_END + (EPS_START - EPS_END) * \
      math.exp(-1. * steps_done / EPS_DECAY)
      if random.random() < eps_threshold:
      action = env.action_space.sample()
      else:
      state = torch.FloatTensor(np.float32(state)).unsqueeze(0).to(device)
      q_values = model(state)
      _, action = q_values.max(1)
      action = int(action.item())
      
      # perform the selected action in the environment
      next_state, reward, done, _ = env.step(action)
      
      # store the transition in the replay buffer
      exp = (state, action, next_state, reward, done)
      replay_buffer.append(exp)
          
      # update the state and step count
      state = next_state
      steps_done += 1
      
      # train the model after every C steps
      if len(replay_buffer) > BATCH_SIZE and steps_done % C == 0:
      for i in range(TRAIN_FREQ):
          batch = random.sample(replay_buffer, BATCH_SIZE)
          train(model, optimizer, loss_fn, batch)
          
              
      episode_rewards.append(reward)
      if done or steps_done >= MAX_STEPS:
      break
      
    return sum(episode_rewards)
      
      
    if __name__ == '__main__':
    NUM_EPISODES = 200
    REWARDS = []
    # initialize target network to same parameters as online network
    target_net = DQN(env.observation_space.shape[0], env.action_space.n)
    target_net.load_state_dict(model.state_dict())
      
    for ep in range(NUM_EPISODES):
      rewards = run_episode()
      REWARDS.append(rewards)
      
      print('[Episode {}/{}] Reward {}'.format(ep+1, NUM_EPISODES, rewards))
      
      if ep % TARGET_UPDATE == 0:
      target_net.load_state_dict(model.state_dict())
          
    # plot the total reward per episode
    plt.plot(REWARDS)
    plt.xlabel('Episode')
    plt.ylabel('Total Reward')
    plt.show()
    
    
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
    
    代码解读

    import os
    import time
    import torch
    import torch.nn as nn
    import gym
    from tensorboardX import SummaryWriter
    import numpy as np
    import math
    class ActorCritic(nn.Module):
    """
    This class implements the ACTOR CRITIC NETWORK used by the A2C algorithm. It takes as input 
    the size of the observations space and outputs two vectors of length n_actions representing
    the probability distribution over possible actions and the expected value of each action respectively.
    """
    def __init__(self, observation_size, hidden_size, n_actions):
      super().__init__()
      self.hidden_size = hidden_size
      self.actor = nn.Sequential(
              nn.Linear(observation_size, hidden_size),
              nn.Tanh(),
              nn.Linear(hidden_size, hidden_size),
              nn.Tanh(),
              nn.Linear(hidden_size, n_actions)
          )
      self.critic = nn.Sequential(
              nn.Linear(observation_size, hidden_size),
              nn.Tanh(),
              nn.Linear(hidden_size, hidden_size),
              nn.Tanh(),
              nn.Linear(hidden_size, 1)
          )
      
    def forward(self, x):
      """
      Forward pass through both actor and critic networks. Returns tuple consisting of 
      the actor output (action probabilities) and the critic output (expected value of each action).
      """
      probs = F.softmax(self.actor(x), dim=-1)
      value = self.critic(x)
      return probs, value
    class A2CAgent:
    """This is a single agent that interacts with the environment."""
    def __init__(self, name, obs_size, act_size, gamma, lr, entropy_coef, max_steps, log_interval, seed):
      self.name = name
      self.obs_size = obs_size
      self.act_size = act_size
      self.gamma = gamma
      self.lr = lr
      self.entropy_coef = entropy_coef
      self.max_steps = max_steps
      self.log_interval = log_interval
      self.seed = seed
      self.training_mode = False
    def choose_action(self, obs, training=True):
      """Choose an action given an observation"""
      self.training_mode = training
      self.model.train(training)
      with torch.no_grad():
      obs = torch.tensor(obs, dtype=torch.float32).unsqueeze(0).to(self.device)
      prob = self.model.actor(obs)[0].detach().cpu().numpy()
      dist = Categorical(prob)
      action = dist.sample().item()
      return action
    def learn(self, rollout):
      """Update policy using the given rollout"""
      obs_batch, acts_batch, rews_batch, vals_batch, dones_batch = map(
      lambda x: torch.cat(x, dim=0).to(self.device),
      zip(*rollout)
      )
      advantages = rews_batch - vals_batch[:-1]
      
      probs, vals = self.model(obs_batch)
      m_probs, m_vals = self.model(obs_batch[-1])
      last_val = m_vals.view(-1).item()
      discounted_rews = utils.discount_rewards(rews_batch + [last_val], self.gamma)[:-1]
      
      val_loss = ((discounted_rews - vals)**2).mean()
      entropy_loss = (-(m_probs*torch.log(probs))).sum()
      pol_loss = -(advantages.detach() * torch.log(probs)).mean()
      
      loss = pol_loss + val_loss + (self.entropy_coef * entropy_loss)
      loss.backward()
      nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=0.5)
      self.optimizer.step()
      
      return {'pol_loss': pol_loss.item(), 
          'val_loss': val_loss.item(), 
          'entropy_loss': entropy_loss.item()}
    def init_model(self, env, device='cpu'):
      """Initialize actor-critic neural networks"""
      self.device = device
      self.model = ActorCritic(self.obs_size, HIDDEN_SIZE, self.act_size).to(self.device)
      self.optimizer = optim.Adam(self.model.parameters(), lr=self.lr)
      
      
    class RolloutStorage:
    """Stores rollouts data until they can be used to update a model"""
    def __init__(self, num_steps, num_processes, obs_size, act_size):
      self.observations = torch.zeros(num_steps+1, num_processes, obs_size)
      self.actions = torch.zeros(num_steps, num_processes, 1).long()
      self.rewards = torch.zeros(num_steps, num_processes, 1)
      self.returns = torch.zeros(num_steps+1, num_processes, 1)
      self.masks = torch.ones(num_steps+1, num_processes, 1)
      self.index = 0
      self.num_steps = num_steps
      self.num_processes = num_processes
      self.obs_size = obs_size
      self.act_size = act_size
      
    def insert(self, current_obs, action, reward, mask):
      """Insert new observation into the storage buffer"""
      self.observations[self.index+1].copy_(current_obs)
      self.actions[self.index].copy_(action)
      self.rewards[self.index].copy_(reward)
      self.masks[self.index+1].copy_(mask)
      self.index = (self.index + 1) % self.num_steps
      
    def after_update(self):
      """Compute returns and clear out the buffers"""
      self._compute_returns()
      self.observations.zero_()
      self.actions.zero_()
      self.rewards.zero_()
      self.masks.fill_(1)
      self.index = 0
      
    def compute_advantages(self, last_val=0.0):
      """Computes advantage estimates based on the current returns"""
      advs = self.returns[:-1] - self.values + last_val
      advs = (advs - advs.mean())/(advs.std()+1e-8)
      return advs
    def feed_forward_generator(self, advantages, mini_batch_size):
      """Generates batches of data from stored rollout"""
      batch_size = self.num_processes * self.num_steps
      assert batch_size >= mini_batch_size, "Batch size should be greater than or equal to sample size"
      indices = torch.randperm(batch_size).tolist()
      for start_idx in range(0, batch_size, mini_batch_size):
      end_idx = start_idx + mini_batch_size
      sampled_indices = indices[start_idx:end_idx]
      yield self._get_samples(sampled_indices, advantages)
      
      
    def _get_samples(self, indices, advantages):
      """Retrieves samples according to the specified indices"""
      obs_batch = self.observations[:-1].view(-1, *self.obs_size)[indices]
      act_batch = self.actions.view(-1, 1)[indices]
      ret_batch = self.returns[:-1].view(-1, 1)[indices]
      adv_batch = advantages.view(-1, 1)[indices]
      old_v_batch = self.values.view(-1, 1)[indices]
      old_p_batch = self.old_log_probs.view(-1, 1)[indices]
      return obs_batch, act_batch, ret_batch, adv_batch, old_v_batch, old_p_batch
    def _compute_returns(self):
      """Computes returns recursively from the rewards"""
      R = 0
      self.returns[-1] = self.rewards
      for t in reversed(range(self.rewards.size(0))):
      R = self.gamma * R + self.rewards[t]
      self.returns[t] = R
    
    
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
    
    代码解读