From 67f7f3500c37f662494315ecd1ec65553b41ac15 Mon Sep 17 00:00:00 2001 From: romue Date: Wed, 16 Jun 2021 17:42:13 +0200 Subject: [PATCH] added own dqn --- algorithms/_base.py | 166 ++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 166 insertions(+) create mode 100644 algorithms/_base.py diff --git a/algorithms/_base.py b/algorithms/_base.py new file mode 100644 index 0000000..5dfc1b7 --- /dev/null +++ b/algorithms/_base.py @@ -0,0 +1,166 @@ +from typing import Tuple, NamedTuple +from collections import namedtuple, deque +import numpy as np +import random +import torch +import torch.nn as nn +import torch.nn.functional as F +from stable_baselines3.common.utils import polyak_update + + +class Experience(NamedTuple): + observation: np.ndarray + next_observation: np.ndarray + action: int + reward: float + done : bool + priority: float = 1 + info : dict = {} + + +class BaseBuffer: + def __init__(self, size: int): + self.size = size + self.experience = deque(maxlen=size) + + def __len__(self): + return len(self.experience) + + def add(self, experience): + self.experience.append(experience) + + def sample(self, k): + sample = random.choices(self.experience, k=k) + return sample + + +class PERBuffer(BaseBuffer): + def __init__(self, size, alpha=0.2): + super(PERBuffer, self).__init__(size) + self.alpha = alpha + + def sample(self, k): + pr = [abs(e.priority)**self.alpha for e in self.experience] + pr = np.array(pr) / sum(pr) + idxs = random.choices(range(len(self)), weights=pr, k=k) + pass + + +class BaseDQN(nn.Module): + def __init__(self): + super(BaseDQN, self).__init__() + self.net = nn.Sequential( + nn.Linear(5 * 5 * 3, 64), + nn.Tanh(), + nn.Linear(64, 64), + nn.Tanh(), + nn.Linear(64, 5) + ) + + def act(self, x): + with torch.no_grad(): + action = self.net(x.view(x.shape[0], -1)).argmax(-1).item() + return action + + def forward(self, x): + return self.net(x.view(x.shape[0], -1)) + + def random_action(self): + return random.randrange(0, 5) + + +class BaseQlearner: + def __init__(self, q_net, target_q_net, env, buffer, n_steps, target_update, warmup, eps_end, + gamma=0.99, train_every_n_steps=4, exploration_fraction=0.2, batch_size=64, lr=1e-4, reg_weight=0.0): + self.q_net = q_net + self.target_q_net = target_q_net + self.target_q_net.eval() + self.env = env + self.buffer = buffer + self.n_steps = n_steps + self.target_update = target_update + self.warmup = warmup + self.eps = 1. + self.eps_end = eps_end + self.exploration_fraction = exploration_fraction + self.batch_size = batch_size + self.gamma = gamma + self.train_every_n_steps = train_every_n_steps + self.lr = lr + self.reg_weight = reg_weight + self.optimizer = torch.optim.Adam(self.q_net.parameters(), lr=self.lr) + self.device = 'cpu' + self.running_reward = deque(maxlen=10) + + def to(self, device): + self.device = device + return self + + def anneal_eps(self, step): + fraction = min(float(step) / int(self.exploration_fraction*self.n_steps), 1.0) + self.eps = 1 + fraction * (self.eps_end - 1) + + def learn(self): + step = 0 + while step < self.n_steps: + obs, done = self.env.reset(), False + total_reward = 0 + while not done: + action = self.q_net.random_action() if np.random.rand() < self.eps \ + else self.q_net.act(torch.from_numpy(obs).unsqueeze(0).float()) + next_obs, reward, done, info = env.step(action) + print(action, reward) + experience = Experience(obs.copy(), next_obs.copy(), action, reward, done) # do we really need to copy? + obs = next_obs + self.buffer.add(experience) + + # end of step routine + self.anneal_eps(step) + step += 1 + total_reward += reward + if step % self.train_every_n_steps == 0: + self.train() + if step % self.target_update == 0: + polyak_update(self.q_net.parameters(), self.target_q_net.parameters(), 1.0) + + self.running_reward.append(total_reward) + if step % 800 == 0: + print(f'Step: {step} ({(step/self.n_steps)*100:.2f}%)\tRunning reward: {sum(list(self.running_reward))/len(self.running_reward)}\t eps: {self.eps:.4f}') + + def train(self): + for _ in range(4): + experience = self.buffer.sample(self.batch_size) + obs = torch.stack([torch.from_numpy(e.observation) for e in experience], 0).float() + next_obs = torch.stack([torch.from_numpy(e.next_observation) for e in experience], 0).float() + actions = torch.tensor([e.action for e in experience]).long() + rewards = torch.tensor([e.reward for e in experience]).float() + dones = torch.tensor([e.done for e in experience]).float() + + with torch.no_grad(): + next_q_values = self.target_q_net(next_obs).max(-1)[0] + target_q_values = rewards + (1 - dones) * self.gamma * next_q_values + + q_values = self.q_net(obs).gather(-1, actions.unsqueeze(-1)) + + delta = q_values - target_q_values + loss = torch.mean(self.reg_weight * q_values + torch.pow(delta, 2)) + + # Optimize the model + self.optimizer.zero_grad() + loss.backward() + torch.nn.utils.clip_grad_norm_(self.q_net.parameters(), 10) + self.optimizer.step() + + +if __name__ == '__main__': + from environments.factory.simple_factory import SimpleFactory, DirtProperties, MovementProperties + dirt_props = DirtProperties(clean_amount=3, gain_amount=0.2, max_global_amount=30, + max_local_amount=5, spawn_frequency=1, max_spawn_ratio=0.05) + move_props = MovementProperties(allow_diagonal_movement=True, + allow_square_movement=True, + allow_no_op=False) + env = SimpleFactory(dirt_properties=dirt_props, movement_properties=move_props, n_agents=1, pomdp_radius=2, combin_agent_slices_in_obs=True, max_steps=400, omit_agent_slice_in_obs=False) + #print(env.action_space) + dqn, target_dqn = BaseDQN(), BaseDQN() + learner = BaseQlearner(dqn, target_dqn, env, BaseBuffer(50000), n_steps=100000, target_update=2000, warmup=1000, train_every_n_steps=1, eps_end=0.025, reg_weight=0.0, exploration_fraction=0.3) + print(learner.learn())