mirror of
https://github.com/Tencent/GameAISDK.git
synced 2026-10-20 14:16:41 +08:00
changqingai
150 lines
6.7 KiB
Python
150 lines
6.7 KiB
Python
# -*- 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()
|