# -*- coding: utf-8 -*- """ Tencent is pleased to support the open source community by making GameAISDK available. This source code file is licensed under the GNU General Public License Version 3. For full details, please refer to the file "LICENSE.txt" which is provided as part of this source code package. Copyright (C) 2020 THL A29 Limited, a Tencent company. All rights reserved. """ from __future__ import division import os import numpy as np import torch from torch import optim from .model import DQN class Agent: def __init__(self, args, action_space): self.action_space = action_space self.args = args #self.atoms = args.atoms #self.Vmin = args.V_min #self.Vmax = args.V_max self.support = torch.linspace(args.V_min, args.V_max, args.atoms).to(device=args.device) # Support (range) of z self.delta_z = (args.V_max - args.V_min) / (args.atoms - 1) #self.batch_size = args.batch_size #self.n = args.multi_step #self.discount = args.discount self.online_net = DQN(args, self.action_space).to(device=args.device) # Load pretrained model if provided if args.model: if os.path.isfile(args.model): # Always load tensors onto CPU by default, will shift to GPU if necessary state_dict = torch.load(args.model, map_location='cpu') if 'conv1.weight' in state_dict.keys(): for old_key, new_key in (('conv1.weight', 'convs.0.weight'), ('conv1.bias', 'convs.0.bias'), ('conv2.weight', 'convs.2.weight'), ('conv2.bias', 'convs.2.bias'), ('conv3.weight', 'convs.4.weight'), ('conv3.bias', 'convs.4.bias')): state_dict[new_key] = state_dict[old_key] # Re-map state dict for old pretrained models del state_dict[old_key] # Delete old keys for strict load_state_dict self.online_net.load_state_dict(state_dict) print("Loading pretrained model: " + args.model) else: # Raise error if incorrect model path provided raise FileNotFoundError(args.model) self.online_net.train() self.target_net = DQN(args, self.action_space).to(device=args.device) self.update_target_net() self.target_net.train() for param in self.target_net.parameters(): param.requires_grad = False self.optimiser = optim.Adam(self.online_net.parameters(), lr=args.learning_rate, eps=args.adam_eps) # Resets noisy weights in all linear layers (of online net only) def reset_noise(self): self.online_net.reset_noise() # Acts based on single state (no batch) def act(self, state): with torch.no_grad(): return (self.online_net(state.unsqueeze(0)) * self.support).sum(2).argmax(1).item() # Acts with an ε-greedy policy (used for evaluation only) def act_e_greedy(self, state, epsilon=0.001): # High ε can reduce evaluation scores drastically return np.random.randint(0, self.action_space) if np.random.random() < epsilon else self.act(state) def learn(self, mem): # Sample transitions 样本采样数据 idxs, states, actions, returns, next_states, non_terminals, weights = mem.sample(self.args.batch_size) # Calculate current state probabilities (online network noise already sampled) # 当前状态计算结果 log_ps = self.online_net(states, log=True) # Log probabilities log p(s_t, ·; θonline) log_ps_a = log_ps[range(self.args.batch_size), actions] # log p(s_t, a_t; θonline) with torch.no_grad(): # Calculate nth next state probabilities pns = self.online_net(next_states) # Probabilities p(s_t+n, ·; θonline) dns = self.support.expand_as(pns) * pns # Distribution d_t+n = (z, p(s_t+n, ·; θonline)) # Perform argmax action selection using online network: argmax_a[(z, p(s_t+n, a; θonline))] arg_max_indices_ns = dns.sum(2).argmax(1) self.target_net.reset_noise() # Sample new target net noise # Probabilities p(s_t+n, ·; θtarget) pns = self.target_net(next_states) # Double-Q probabilities p(s_t+n, argmax_a[(z, p(s_t+n, a; θonline))]; θtarget) pns_a = pns[range(self.args.batch_size), arg_max_indices_ns] # Compute Tz (Bellman operator T applied to z) tz = returns.unsqueeze(1) + non_terminals * (self.args.discount ** self.args.multi_step) * \ self.support.unsqueeze(0) # Clamp between supported values tz = tz.clamp(min=self.args.V_min, max=self.args.V_max) # Compute L2 projection of Tz onto fixed support z b = (tz - self.args.V_min) / self.delta_z l, u = b.floor().to(torch.int64), b.ceil().to(torch.int64) # Fix disappearing probability mass when l = b = u (b is int) l[(u > 0) * (l == u)] -= 1 u[(l < (self.args.atoms - 1)) * (l == u)] += 1 # Distribute probability of Tz m = states.new_zeros(self.args.batch_size, self.args.atoms) offset = torch.linspace(0, ((self.args.batch_size - 1) * self.args.atoms), self.args.batch_size).unsqueeze(1).expand( self.args.batch_size, self.args.atoms).to(actions) m.view(-1).index_add_(0, (l + offset).view(-1), (pns_a * (u.float() - b)).view(-1)) # m_l = m_l + p(s_t+n, a*)(u - b) m.view(-1).index_add_(0, (u + offset).view(-1), (pns_a * (b - l.float())).view(-1)) # m_u = m_u + p(s_t+n, a*)(b - l) # 交叉熵损失函数 loss = -torch.sum(m * log_ps_a, 1) # Cross-entropy loss (minimises DKL(m||p(s_t, a_t))) self.online_net.zero_grad() (weights * loss).mean().backward() # Back-propagate importance-weighted mini-batch loss self.optimiser.step() mem.update_priorities(idxs, loss.detach().cpu().numpy()) # Update priorities of sampled transitions def update_target_net(self): self.target_net.load_state_dict(self.online_net.state_dict()) # Save model parameters on current device (don't move model between devices) def save(self, path, name='model.pth'): torch.save(self.online_net.state_dict(), os.path.join(path, name)) # Evaluates Q-value based on single state (no batch) def evaluate_q(self, state): with torch.no_grad(): return (self.online_net(state.unsqueeze(0)) * self.support).sum(2).max(1)[0].item() def train(self): self.online_net.train() def eval(self): self.online_net.eval()