From e541e3427036bf5a21b1c8c5800f6bfff9dd8581 Mon Sep 17 00:00:00 2001 From: romue Date: Fri, 18 Jun 2021 13:55:38 +0200 Subject: [PATCH 1/9] add CER sampling and Munchhausen DQN --- algorithms/_base.py | 112 ++++++++++++++++++++++++++++++++------------ 1 file changed, 81 insertions(+), 31 deletions(-) diff --git a/algorithms/_base.py b/algorithms/_base.py index 8ca0aef..bbb2b4a 100644 --- a/algorithms/_base.py +++ b/algorithms/_base.py @@ -1,4 +1,4 @@ -from typing import NamedTuple, Union +from typing import NamedTuple, Union, Iterable from collections import namedtuple, deque import numpy as np import random @@ -30,8 +30,9 @@ class BaseBuffer: def add(self, experience): self.experience.append(experience) - def sample(self, k): - sample = random.choices(self.experience, k=k) + def sample(self, k, cer=4): + sample = random.choices(self.experience, k=k-cer) + for i in range(cer): sample += [self.experience[-i]] observations = torch.stack([torch.from_numpy(e.observation) for e in sample], 0).float() next_observations = torch.stack([torch.from_numpy(e.next_observation) for e in sample], 0).float() actions = torch.tensor([e.action for e in sample]).long() @@ -40,18 +41,6 @@ class BaseBuffer: return Experience(observations, next_observations, actions, rewards, dones) -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__() @@ -80,14 +69,21 @@ class BaseDQN(nn.Module): return random.randrange(0, 5) +def soft_update(local_model, target_model, tau): + # taken from https://github.com/BY571/Munchausen-RL/blob/master/M-DQN.ipynb + for target_param, local_param in zip(target_model.parameters(), local_model.parameters()): + target_param.data.copy_(tau*local_param.data + (1.-tau)*target_param.data) + + class BaseQlearner: def __init__(self, q_net, target_q_net, env, buffer, target_update, eps_end, n_agents=1, - gamma=0.99, train_every_n_steps=4, n_grad_steps=1, + gamma=0.99, train_every_n_steps=4, n_grad_steps=1, tau=1.0, max_grad_norm=10, 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.q_net.apply(self.weights_init) + #self.q_net.apply(self.weights_init) self.target_q_net.eval() + soft_update(self.q_net, self.target_q_net, tau=1.0) self.env = env self.buffer = buffer self.target_update = target_update @@ -99,10 +95,12 @@ class BaseQlearner: self.train_every_n_steps = train_every_n_steps self.n_grad_steps = n_grad_steps self.lr = lr + self.tau = tau self.reg_weight = reg_weight self.n_agents = n_agents self.device = 'cpu' self.optimizer = torch.optim.AdamW(self.q_net.parameters(), lr=self.lr) + self.max_grad_norm = max_grad_norm self.running_reward = deque(maxlen=5) self.running_loss = deque(maxlen=5) self._n_updates = 0 @@ -112,7 +110,7 @@ class BaseQlearner: return self @staticmethod - def weights_init(module, activation='relu'): + def weights_init(module, activation='leaky_relu'): if isinstance(module, (nn.Linear, nn.Conv2d)): nn.init.xavier_normal_(module.weight, gain=torch.nn.init.calculate_gain(activation)) if module.bias is not None: @@ -154,35 +152,38 @@ class BaseQlearner: self._n_updates += 1 if step % self.target_update == 0: print('UPDATE') - polyak_update(self.q_net.parameters(), self.target_q_net.parameters(), 1) - + soft_update(self.q_net, self.target_q_net, tau=self.tau) self.running_reward.append(total_reward) if step % 10 == 0: print(f'Step: {step} ({(step/n_steps)*100:.2f}%)\tRunning reward: {sum(list(self.running_reward))/len(self.running_reward):.2f}\t' f' eps: {self.eps:.4f}\tRunning loss: {sum(list(self.running_loss))/len(self.running_loss):.4f}\tUpdates:{self._n_updates}') - def _training_routine(self, obs, next_obs, action): + def _training_routine(self, obs, next_obs, action, reward): current_q_values = self.q_net(obs) current_q_values = torch.gather(current_q_values, dim=-1, index=action) next_q_values_raw = self.target_q_net(next_obs).max(dim=-1)[0].reshape(-1, 1).detach() return current_q_values, next_q_values_raw + def train(self): if len(self.buffer) < self.batch_size: return for _ in range(self.n_grad_steps): - experience = self.buffer.sample(self.batch_size) - #print(experience.observation.shape, experience.next_observation.shape, experience.action.shape, experience.reward.shape, experience.done.shape) + experience = self.buffer.sample(self.batch_size, cer=self.train_every_n_steps) + if self.n_agents <= 1: - pred_q, target_q_raw = self._training_routine(experience.observation, experience.next_observation, experience.action) + pred_q, target_q_raw = self._training_routine(experience.observation, + experience.next_observation, + experience.action, + experience.reward) else: - pred_q, target_q_raw = torch.zeros((self.batch_size, 1)), torch.zeros((self.batch_size, 1)) + pred_q, target_q_raw, reward = [torch.zeros((self.batch_size, 1))]*3 for agent_i in range(self.n_agents): q_values, next_q_values_raw = self._training_routine(experience.observation[:, agent_i], - experience.next_observation[:, agent_i], - experience.action[:, agent_i].unsqueeze(-1) - ) + experience.next_observation[:, agent_i], + experience.action[:, agent_i].unsqueeze(-1), + experience.reward) pred_q += q_values target_q_raw += next_q_values_raw target_q = experience.reward + (1 - experience.done) * self.gamma * target_q_raw @@ -193,7 +194,56 @@ class BaseQlearner: # Optimize the model self.optimizer.zero_grad() loss.backward() - torch.nn.utils.clip_grad_norm_(self.q_net.parameters(), 10) + torch.nn.utils.clip_grad_norm_(self.q_net.parameters(), self.max_grad_norm) + self.optimizer.step() + + +class MDQN(BaseQlearner): + def __init__(self, *args, temperature=0.03, alpha=0.9, clip_l0=-1.0, **kwargs): + super(MDQN, self).__init__(*args, **kwargs) + assert self.n_agents == 1, 'M-DQN currently only supports single agent training' + self.temperature = temperature + self.alpha = alpha + self.clip0 = clip_l0 + + def train(self): + if len(self.buffer) < self.batch_size: return + for _ in range(self.n_grad_steps): + + experience = self.buffer.sample(self.batch_size, cer=self.train_every_n_steps) + + q_target_next = self.target_q_net(experience.next_observation).detach() + advantages_next = (q_target_next - q_target_next.max(-1)[0].unsqueeze(-1)) + logsum = torch.logsumexp(advantages_next / self.temperature, -1).unsqueeze(-1) + tau_log_pi_next = advantages_next - self.temperature * logsum + + pi_target = F.softmax(q_target_next / self.temperature, dim=-1) + q_target = (self.gamma * (pi_target * (q_target_next - tau_log_pi_next) * (1 - experience.done)).sum(-1)).unsqueeze(-1) + + q_k_targets = self.target_q_net(experience.observation).detach() + v_k_target = q_k_targets.max(-1)[0].unsqueeze(-1) + logsum = torch.logsumexp((q_k_targets - v_k_target) / self.temperature, -1).unsqueeze(-1) + log_pi = q_k_targets - v_k_target - self.temperature * logsum + munchausen_addon = log_pi.gather(-1, experience.action) + + munchausen_reward = (experience.reward + self.alpha * torch.clamp(munchausen_addon, min=self.clip0, max=0)) + + # Compute Q targets for current states + m_q_target = munchausen_reward + q_target + + # Get expected Q values from local model + q_k = self.q_net(experience.observation) + pred_q = q_k.gather(-1, experience.action) + + # Compute loss + loss = torch.mean(self.reg_weight * pred_q + torch.pow(pred_q - m_q_target, 2)) + + # log loss + self.running_loss.append(loss.item()) + # Optimize the model + self.optimizer.zero_grad() + loss.backward() + torch.nn.utils.clip_grad_norm_(self.q_net.parameters(), self.max_grad_norm) self.optimizer.step() @@ -221,6 +271,6 @@ if __name__ == '__main__': dqn, target_dqn = BaseDQN(), BaseDQN() - learner = BaseQlearner(dqn, target_dqn, env, BaseBuffer(40000), target_update=3500, lr=0.0008, gamma=0.99, n_agents=N_AGENTS, - train_every_n_steps=4, eps_end=0.05, n_grad_steps=1, reg_weight=0.05, exploration_fraction=0.25, batch_size=64) + learner = MDQN(dqn, target_dqn, env, BaseBuffer(40000), target_update=3500, lr=0.0008, gamma=0.99, n_agents=N_AGENTS, tau=0.95, max_grad_norm=10, + train_every_n_steps=4, eps_end=0.025, n_grad_steps=1, reg_weight=0.1, exploration_fraction=0.25, batch_size=64) learner.learn(100000) From bd63c603eeb4dc02ca5f766dda5717fa91270a6c Mon Sep 17 00:00:00 2001 From: romue Date: Mon, 21 Jun 2021 10:42:35 +0200 Subject: [PATCH 2/9] add Munchhausen DQN refactoring --- algorithms/_base.py | 21 +++++++++++++-------- 1 file changed, 13 insertions(+), 8 deletions(-) diff --git a/algorithms/_base.py b/algorithms/_base.py index bbb2b4a..b9d7d83 100644 --- a/algorithms/_base.py +++ b/algorithms/_base.py @@ -206,6 +206,14 @@ class MDQN(BaseQlearner): self.alpha = alpha self.clip0 = clip_l0 + def tau_ln_pi(self, qs): + # Custom log-sum-exp trick from page 18 to compute the e log-policy terms + v_k = qs.max(-1)[0].unsqueeze(-1) + advantage = qs - v_k + logsum = torch.logsumexp(advantage / self.temperature, -1).unsqueeze(-1) + tau_ln_pi = advantage - self.temperature * logsum + return tau_ln_pi + def train(self): if len(self.buffer) < self.batch_size: return for _ in range(self.n_grad_steps): @@ -213,17 +221,14 @@ class MDQN(BaseQlearner): experience = self.buffer.sample(self.batch_size, cer=self.train_every_n_steps) q_target_next = self.target_q_net(experience.next_observation).detach() - advantages_next = (q_target_next - q_target_next.max(-1)[0].unsqueeze(-1)) - logsum = torch.logsumexp(advantages_next / self.temperature, -1).unsqueeze(-1) - tau_log_pi_next = advantages_next - self.temperature * logsum + tau_log_pi_next = self.tau_ln_pi(q_target_next) + + q_k_targets = self.target_q_net(experience.observation).detach() + log_pi = self.tau_ln_pi(q_k_targets) pi_target = F.softmax(q_target_next / self.temperature, dim=-1) q_target = (self.gamma * (pi_target * (q_target_next - tau_log_pi_next) * (1 - experience.done)).sum(-1)).unsqueeze(-1) - q_k_targets = self.target_q_net(experience.observation).detach() - v_k_target = q_k_targets.max(-1)[0].unsqueeze(-1) - logsum = torch.logsumexp((q_k_targets - v_k_target) / self.temperature, -1).unsqueeze(-1) - log_pi = q_k_targets - v_k_target - self.temperature * logsum munchausen_addon = log_pi.gather(-1, experience.action) munchausen_reward = (experience.reward + self.alpha * torch.clamp(munchausen_addon, min=self.clip0, max=0)) @@ -272,5 +277,5 @@ if __name__ == '__main__': dqn, target_dqn = BaseDQN(), BaseDQN() learner = MDQN(dqn, target_dqn, env, BaseBuffer(40000), target_update=3500, lr=0.0008, gamma=0.99, n_agents=N_AGENTS, tau=0.95, max_grad_norm=10, - train_every_n_steps=4, eps_end=0.025, n_grad_steps=1, reg_weight=0.1, exploration_fraction=0.25, batch_size=64) + train_every_n_steps=4, eps_end=0.025, n_grad_steps=1, reg_weight=0.1, exploration_fraction=0.25, batch_size=64) learner.learn(100000) From 543c2987e024e8253d2958c2865cdf151d2969e9 Mon Sep 17 00:00:00 2001 From: romue Date: Mon, 21 Jun 2021 14:05:20 +0200 Subject: [PATCH 3/9] add Munchhausen DQN refactoring --- algorithms/_base.py | 55 ++++++++++++++------------------------------- 1 file changed, 17 insertions(+), 38 deletions(-) diff --git a/algorithms/_base.py b/algorithms/_base.py index b9d7d83..520c690 100644 --- a/algorithms/_base.py +++ b/algorithms/_base.py @@ -1,22 +1,20 @@ -from typing import NamedTuple, Union, Iterable -from collections import namedtuple, deque +from typing import NamedTuple, Union +from collections import 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 -from stable_baselines3.common.buffers import ReplayBuffer -import copy class Experience(NamedTuple): + # can be use for a single (s_t, a, r s_{t+1}) tuple + # or for a batch of tuples observation: np.ndarray next_observation: np.ndarray action: np.ndarray reward: Union[float, np.ndarray] done : Union[bool, np.ndarray] - priority: np.ndarray = 1 class BaseBuffer: @@ -65,9 +63,6 @@ class BaseDQN(nn.Module): values = self.value_head(features) return values + (advantages - advantages.mean()) - def random_action(self): - return random.randrange(0, 5) - def soft_update(local_model, target_model, tau): # taken from https://github.com/BY571/Munchausen-RL/blob/master/M-DQN.ipynb @@ -165,6 +160,14 @@ class BaseQlearner: next_q_values_raw = self.target_q_net(next_obs).max(dim=-1)[0].reshape(-1, 1).detach() return current_q_values, next_q_values_raw + def _backprop_loss(self, loss): + # log loss + self.running_loss.append(loss.item()) + # Optimize the model + self.optimizer.zero_grad() + loss.backward() + torch.nn.utils.clip_grad_norm_(self.q_net.parameters(), self.max_grad_norm) + self.optimizer.step() def train(self): if len(self.buffer) < self.batch_size: return @@ -188,14 +191,7 @@ class BaseQlearner: target_q_raw += next_q_values_raw target_q = experience.reward + (1 - experience.done) * self.gamma * target_q_raw loss = torch.mean(self.reg_weight * pred_q + torch.pow(pred_q - target_q, 2)) - - # log loss - self.running_loss.append(loss.item()) - # Optimize the model - self.optimizer.zero_grad() - loss.backward() - torch.nn.utils.clip_grad_norm_(self.q_net.parameters(), self.max_grad_norm) - self.optimizer.step() + self._backprop_loss(loss) class MDQN(BaseQlearner): @@ -207,7 +203,8 @@ class MDQN(BaseQlearner): self.clip0 = clip_l0 def tau_ln_pi(self, qs): - # Custom log-sum-exp trick from page 18 to compute the e log-policy terms + # computes log(softmax(qs/temperature)) + # Custom log-sum-exp trick from page 18 to compute the log-policy terms v_k = qs.max(-1)[0].unsqueeze(-1) advantage = qs - v_k logsum = torch.logsumexp(advantage / self.temperature, -1).unsqueeze(-1) @@ -242,21 +239,11 @@ class MDQN(BaseQlearner): # Compute loss loss = torch.mean(self.reg_weight * pred_q + torch.pow(pred_q - m_q_target, 2)) - - # log loss - self.running_loss.append(loss.item()) - # Optimize the model - self.optimizer.zero_grad() - loss.backward() - torch.nn.utils.clip_grad_norm_(self.q_net.parameters(), self.max_grad_norm) - self.optimizer.step() - + self._backprop_loss(loss) if __name__ == '__main__': from environments.factory.simple_factory import SimpleFactory, DirtProperties, MovementProperties - from algorithms.reg_dqn import RegDQN - from stable_baselines3.common.vec_env import DummyVecEnv N_AGENTS = 1 @@ -266,16 +253,8 @@ if __name__ == '__main__': allow_square_movement=True, allow_no_op=False) env = SimpleFactory(dirt_properties=dirt_props, movement_properties=move_props, n_agents=N_AGENTS, pomdp_radius=2, max_steps=400, omit_agent_slice_in_obs=False, combin_agent_slices_in_obs=True) - #env = DummyVecEnv([lambda: env]) - from stable_baselines3.dqn import DQN - - #dqn = RegDQN('MlpPolicy', env, verbose=True, buffer_size = 40000, learning_starts = 0, batch_size = 64,learning_rate=0.0008, - # target_update_interval = 3500, exploration_fraction = 0.25, exploration_final_eps = 0.05, - # train_freq=4, gradient_steps=1, reg_weight=0.05, seed=69) - #dqn.learn(100000) - dqn, target_dqn = BaseDQN(), BaseDQN() - learner = MDQN(dqn, target_dqn, env, BaseBuffer(40000), target_update=3500, lr=0.0008, gamma=0.99, n_agents=N_AGENTS, tau=0.95, max_grad_norm=10, + learner = MDQN(dqn, target_dqn, env, BaseBuffer(40000), target_update=3500, lr=0.0007, gamma=0.99, n_agents=N_AGENTS, tau=0.95, max_grad_norm=10, train_every_n_steps=4, eps_end=0.025, n_grad_steps=1, reg_weight=0.1, exploration_fraction=0.25, batch_size=64) learner.learn(100000) From c5d677e9ba791cf1cddac7f57a6a24afebb1fe1d Mon Sep 17 00:00:00 2001 From: romue Date: Tue, 22 Jun 2021 14:19:13 +0200 Subject: [PATCH 4/9] add VDN fix --- algorithms/_base.py | 47 ++++++++++++++++++++++++++++++++------------- 1 file changed, 34 insertions(+), 13 deletions(-) diff --git a/algorithms/_base.py b/algorithms/_base.py index 520c690..1aa28fd 100644 --- a/algorithms/_base.py +++ b/algorithms/_base.py @@ -39,9 +39,9 @@ class BaseBuffer: return Experience(observations, next_observations, actions, rewards, dones) -class BaseDQN(nn.Module): +class BaseDDQN(nn.Module): def __init__(self): - super(BaseDQN, self).__init__() + super(BaseDDQN, self).__init__() self.net = nn.Sequential( nn.Flatten(), nn.Linear(3*5*5, 64), @@ -64,6 +64,27 @@ class BaseDQN(nn.Module): return values + (advantages - advantages.mean()) +class BaseDQN(nn.Module): + def __init__(self): + super(BaseDQN, self).__init__() + self.net = nn.Sequential( + nn.Flatten(), + nn.Linear(3*5*5, 64), + nn.ELU(), + nn.Linear(64, 64), + nn.ELU(), + nn.Linear(64, 9) + ) + + def act(self, x) -> np.ndarray: + with torch.no_grad(): + action = self.forward(x).max(-1)[1].numpy() + return action + + def forward(self, x): + return self.net(x) + + def soft_update(local_model, target_model, tau): # taken from https://github.com/BY571/Munchausen-RL/blob/master/M-DQN.ipynb for target_param, local_param in zip(target_model.parameters(), local_model.parameters()): @@ -154,10 +175,11 @@ class BaseQlearner: print(f'Step: {step} ({(step/n_steps)*100:.2f}%)\tRunning reward: {sum(list(self.running_reward))/len(self.running_reward):.2f}\t' f' eps: {self.eps:.4f}\tRunning loss: {sum(list(self.running_loss))/len(self.running_loss):.4f}\tUpdates:{self._n_updates}') - def _training_routine(self, obs, next_obs, action, reward): + def _training_routine(self, obs, next_obs, action): current_q_values = self.q_net(obs) current_q_values = torch.gather(current_q_values, dim=-1, index=action) next_q_values_raw = self.target_q_net(next_obs).max(dim=-1)[0].reshape(-1, 1).detach() + #print(current_q_values.shape, next_q_values_raw.shape) return current_q_values, next_q_values_raw def _backprop_loss(self, loss): @@ -174,22 +196,21 @@ class BaseQlearner: for _ in range(self.n_grad_steps): experience = self.buffer.sample(self.batch_size, cer=self.train_every_n_steps) - if self.n_agents <= 1: pred_q, target_q_raw = self._training_routine(experience.observation, - experience.next_observation, - experience.action, - experience.reward) + experience.next_observation, + experience.action) + else: - pred_q, target_q_raw, reward = [torch.zeros((self.batch_size, 1))]*3 + pred_q, target_q_raw = torch.zeros((self.batch_size, 1)), torch.zeros((self.batch_size, 1)) for agent_i in range(self.n_agents): q_values, next_q_values_raw = self._training_routine(experience.observation[:, agent_i], - experience.next_observation[:, agent_i], - experience.action[:, agent_i].unsqueeze(-1), - experience.reward) + experience.next_observation[:, agent_i], + experience.action[:, agent_i].unsqueeze(-1)) pred_q += q_values target_q_raw += next_q_values_raw target_q = experience.reward + (1 - experience.done) * self.gamma * target_q_raw + #print(pred_q[0], target_q_raw[0], target_q[0], experience.reward[0]) loss = torch.mean(self.reg_weight * pred_q + torch.pow(pred_q - target_q, 2)) self._backprop_loss(loss) @@ -245,7 +266,7 @@ class MDQN(BaseQlearner): if __name__ == '__main__': from environments.factory.simple_factory import SimpleFactory, DirtProperties, MovementProperties - N_AGENTS = 1 + N_AGENTS = 2 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) @@ -255,6 +276,6 @@ if __name__ == '__main__': env = SimpleFactory(dirt_properties=dirt_props, movement_properties=move_props, n_agents=N_AGENTS, pomdp_radius=2, max_steps=400, omit_agent_slice_in_obs=False, combin_agent_slices_in_obs=True) dqn, target_dqn = BaseDQN(), BaseDQN() - learner = MDQN(dqn, target_dqn, env, BaseBuffer(40000), target_update=3500, lr=0.0007, gamma=0.99, n_agents=N_AGENTS, tau=0.95, max_grad_norm=10, + learner = BaseQlearner(dqn, target_dqn, env, BaseBuffer(40000), target_update=3500, lr=0.0007, gamma=0.99, n_agents=N_AGENTS, tau=0.95, max_grad_norm=10, train_every_n_steps=4, eps_end=0.025, n_grad_steps=1, reg_weight=0.1, exploration_fraction=0.25, batch_size=64) learner.learn(100000) From b5d729e597f0bc1827931a5a7fecbb5a6f91b96f Mon Sep 17 00:00:00 2001 From: romue Date: Tue, 22 Jun 2021 16:23:39 +0200 Subject: [PATCH 5/9] added mlpmaker --- algorithms/{_base.py => dqn.py} | 82 ++++++++++++++++----------------- 1 file changed, 39 insertions(+), 43 deletions(-) rename algorithms/{_base.py => dqn.py} (88%) diff --git a/algorithms/_base.py b/algorithms/dqn.py similarity index 88% rename from algorithms/_base.py rename to algorithms/dqn.py index 1aa28fd..15bed40 100644 --- a/algorithms/_base.py +++ b/algorithms/dqn.py @@ -1,5 +1,5 @@ from typing import NamedTuple, Union -from collections import deque +from collections import deque, OrderedDict import numpy as np import random import torch @@ -39,42 +39,27 @@ class BaseBuffer: return Experience(observations, next_observations, actions, rewards, dones) -class BaseDDQN(nn.Module): - def __init__(self): - super(BaseDDQN, self).__init__() - self.net = nn.Sequential( - nn.Flatten(), - nn.Linear(3*5*5, 64), - nn.ELU(), - nn.Linear(64, 64), - nn.ELU() - ) - self.value_head = nn.Linear(64, 1) - self.advantage_head = nn.Linear(64, 9) - def act(self, x) -> np.ndarray: - with torch.no_grad(): - action = self.forward(x).max(-1)[1].numpy() - return action - def forward(self, x): - features = self.net(x) - advantages = self.advantage_head(features) - values = self.value_head(features) - return values + (advantages - advantages.mean()) +def soft_update(local_model, target_model, tau): + # taken from https://github.com/BY571/Munchausen-RL/blob/master/M-DQN.ipynb + for target_param, local_param in zip(target_model.parameters(), local_model.parameters()): + target_param.data.copy_(tau*local_param.data + (1.-tau)*target_param.data) + + +def mlp_maker(dims): + layers = [('Flatten', nn.Flatten())] + for i in range(1, len(dims)): + layers.append((f'Linear#{i - 1}', nn.Linear(dims[i - 1], dims[i]))) + if i != len(dims) - 1: + layers.append(('ELU', nn.ELU())) + return nn.Sequential(OrderedDict(layers)) class BaseDQN(nn.Module): - def __init__(self): + def __init__(self, dims=[3*5*5, 64, 64, 9]): super(BaseDQN, self).__init__() - self.net = nn.Sequential( - nn.Flatten(), - nn.Linear(3*5*5, 64), - nn.ELU(), - nn.Linear(64, 64), - nn.ELU(), - nn.Linear(64, 9) - ) + self.net = mlp_maker(dims) def act(self, x) -> np.ndarray: with torch.no_grad(): @@ -85,23 +70,34 @@ class BaseDQN(nn.Module): return self.net(x) -def soft_update(local_model, target_model, tau): - # taken from https://github.com/BY571/Munchausen-RL/blob/master/M-DQN.ipynb - for target_param, local_param in zip(target_model.parameters(), local_model.parameters()): - target_param.data.copy_(tau*local_param.data + (1.-tau)*target_param.data) +class BaseDDQN(BaseDQN): + def __init__(self, + backbone_dims=[3*5*5, 64, 64], + value_dims=[64,1], + advantage_dims=[64,9]): + super(BaseDDQN, self).__init__(backbone_dims) + self.value_head = mlp_maker(value_dims) + self.advantage_head = mlp_maker(advantage_dims) + + def forward(self, x): + features = self.net(x) + advantages = self.advantage_head(features) + values = self.value_head(features) + return values + (advantages - advantages.mean()) + class BaseQlearner: - def __init__(self, q_net, target_q_net, env, buffer, target_update, eps_end, n_agents=1, + def __init__(self, q_net, target_q_net, env, buffer_size, target_update, eps_end, n_agents=1, gamma=0.99, train_every_n_steps=4, n_grad_steps=1, tau=1.0, max_grad_norm=10, exploration_fraction=0.2, batch_size=64, lr=1e-4, reg_weight=0.0): self.q_net = q_net + print(self.q_net) self.target_q_net = target_q_net - #self.q_net.apply(self.weights_init) self.target_q_net.eval() soft_update(self.q_net, self.target_q_net, tau=1.0) self.env = env - self.buffer = buffer + self.buffer = BaseBuffer(buffer_size) self.target_update = target_update self.eps = 1. self.eps_end = eps_end @@ -128,7 +124,7 @@ class BaseQlearner: @staticmethod def weights_init(module, activation='leaky_relu'): if isinstance(module, (nn.Linear, nn.Conv2d)): - nn.init.xavier_normal_(module.weight, gain=torch.nn.init.calculate_gain(activation)) + nn.init.orthogonal_(module.weight, gain=torch.nn.init.calculate_gain(activation)) if module.bias is not None: module.bias.data.fill_(0.0) @@ -179,7 +175,6 @@ class BaseQlearner: current_q_values = self.q_net(obs) current_q_values = torch.gather(current_q_values, dim=-1, index=action) next_q_values_raw = self.target_q_net(next_obs).max(dim=-1)[0].reshape(-1, 1).detach() - #print(current_q_values.shape, next_q_values_raw.shape) return current_q_values, next_q_values_raw def _backprop_loss(self, loss): @@ -265,8 +260,9 @@ class MDQN(BaseQlearner): if __name__ == '__main__': from environments.factory.simple_factory import SimpleFactory, DirtProperties, MovementProperties + from gym.wrappers import FrameStack - N_AGENTS = 2 + N_AGENTS = 1 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) @@ -275,7 +271,7 @@ if __name__ == '__main__': allow_no_op=False) env = SimpleFactory(dirt_properties=dirt_props, movement_properties=move_props, n_agents=N_AGENTS, pomdp_radius=2, max_steps=400, omit_agent_slice_in_obs=False, combin_agent_slices_in_obs=True) - dqn, target_dqn = BaseDQN(), BaseDQN() - learner = BaseQlearner(dqn, target_dqn, env, BaseBuffer(40000), target_update=3500, lr=0.0007, gamma=0.99, n_agents=N_AGENTS, tau=0.95, max_grad_norm=10, + dqn, target_dqn = BaseDDQN(), BaseDDQN() + learner = MDQN(dqn, target_dqn, env, 40000, target_update=3500, lr=0.0007, gamma=0.99, n_agents=N_AGENTS, tau=0.95, max_grad_norm=10, train_every_n_steps=4, eps_end=0.025, n_grad_steps=1, reg_weight=0.1, exploration_fraction=0.25, batch_size=64) learner.learn(100000) From f0bb5071214c901ed35cfffc202342444ba60b2d Mon Sep 17 00:00:00 2001 From: romue Date: Tue, 22 Jun 2021 17:28:22 +0200 Subject: [PATCH 6/9] added mlpmaker --- algorithms/dqn.py | 28 ++++++++++++++-------------- 1 file changed, 14 insertions(+), 14 deletions(-) diff --git a/algorithms/dqn.py b/algorithms/dqn.py index 15bed40..7ce2d20 100644 --- a/algorithms/dqn.py +++ b/algorithms/dqn.py @@ -47,19 +47,22 @@ def soft_update(local_model, target_model, tau): target_param.data.copy_(tau*local_param.data + (1.-tau)*target_param.data) -def mlp_maker(dims): - layers = [('Flatten', nn.Flatten())] +def mlp_maker(dims, flatten=False, activation='elu', activation_last='identity'): + activations = {'elu': nn.ELU, 'relu': nn.ReLU, + 'leaky_relu': nn.LeakyReLU, 'tanh': nn.Tanh, + 'gelu': nn.GELU, 'identity': nn.Identity} + layers = [('Flatten', nn.Flatten())] if flatten else [] for i in range(1, len(dims)): - layers.append((f'Linear#{i - 1}', nn.Linear(dims[i - 1], dims[i]))) - if i != len(dims) - 1: - layers.append(('ELU', nn.ELU())) + layers.append((f'Layer #{i - 1}: Linear', nn.Linear(dims[i - 1], dims[i]))) + activation_str = activation if i != len(dims)-1 else activation_last + layers.append((f'Layer #{i - 1}: {activation_str.capitalize()}', activations[activation_str]())) return nn.Sequential(OrderedDict(layers)) class BaseDQN(nn.Module): def __init__(self, dims=[3*5*5, 64, 64, 9]): super(BaseDQN, self).__init__() - self.net = mlp_maker(dims) + self.net = mlp_maker(dims, flatten=True) def act(self, x) -> np.ndarray: with torch.no_grad(): @@ -76,6 +79,7 @@ class BaseDDQN(BaseDQN): value_dims=[64,1], advantage_dims=[64,9]): super(BaseDDQN, self).__init__(backbone_dims) + self.net = mlp_maker(backbone_dims, flatten=True) self.value_head = mlp_maker(value_dims) self.advantage_head = mlp_maker(advantage_dims) @@ -86,13 +90,11 @@ class BaseDDQN(BaseDQN): return values + (advantages - advantages.mean()) - class BaseQlearner: def __init__(self, q_net, target_q_net, env, buffer_size, target_update, eps_end, n_agents=1, gamma=0.99, train_every_n_steps=4, n_grad_steps=1, tau=1.0, max_grad_norm=10, exploration_fraction=0.2, batch_size=64, lr=1e-4, reg_weight=0.0): self.q_net = q_net - print(self.q_net) self.target_q_net = target_q_net self.target_q_net.eval() soft_update(self.q_net, self.target_q_net, tau=1.0) @@ -205,14 +207,13 @@ class BaseQlearner: pred_q += q_values target_q_raw += next_q_values_raw target_q = experience.reward + (1 - experience.done) * self.gamma * target_q_raw - #print(pred_q[0], target_q_raw[0], target_q[0], experience.reward[0]) loss = torch.mean(self.reg_weight * pred_q + torch.pow(pred_q - target_q, 2)) self._backprop_loss(loss) -class MDQN(BaseQlearner): +class MunchhausenQLearner(BaseQlearner): def __init__(self, *args, temperature=0.03, alpha=0.9, clip_l0=-1.0, **kwargs): - super(MDQN, self).__init__(*args, **kwargs) + super(MunchhausenQLearner, self).__init__(*args, **kwargs) assert self.n_agents == 1, 'M-DQN currently only supports single agent training' self.temperature = temperature self.alpha = alpha @@ -260,7 +261,6 @@ class MDQN(BaseQlearner): if __name__ == '__main__': from environments.factory.simple_factory import SimpleFactory, DirtProperties, MovementProperties - from gym.wrappers import FrameStack N_AGENTS = 1 @@ -272,6 +272,6 @@ if __name__ == '__main__': env = SimpleFactory(dirt_properties=dirt_props, movement_properties=move_props, n_agents=N_AGENTS, pomdp_radius=2, max_steps=400, omit_agent_slice_in_obs=False, combin_agent_slices_in_obs=True) dqn, target_dqn = BaseDDQN(), BaseDDQN() - learner = MDQN(dqn, target_dqn, env, 40000, target_update=3500, lr=0.0007, gamma=0.99, n_agents=N_AGENTS, tau=0.95, max_grad_norm=10, - train_every_n_steps=4, eps_end=0.025, n_grad_steps=1, reg_weight=0.1, exploration_fraction=0.25, batch_size=64) + learner = MunchhausenQLearner(dqn, target_dqn, env, 40000, target_update=3500, lr=0.0007, gamma=0.99, n_agents=N_AGENTS, tau=0.95, max_grad_norm=10, + train_every_n_steps=4, eps_end=0.025, n_grad_steps=1, reg_weight=0.1, exploration_fraction=0.25, batch_size=64) learner.learn(100000) From 42f0dde05643937dd766fd068cb075b48c59a5d3 Mon Sep 17 00:00:00 2001 From: romue Date: Wed, 23 Jun 2021 10:56:18 +0200 Subject: [PATCH 7/9] refactored algorithms --- algorithms/common.py | 102 ++++++++++++++ algorithms/dqn.py | 277 -------------------------------------- algorithms/m_q_learner.py | 53 ++++++++ algorithms/q_learner.py | 144 ++++++++++++++++++++ algorithms/vdn_learner.py | 23 ++++ 5 files changed, 322 insertions(+), 277 deletions(-) create mode 100644 algorithms/common.py delete mode 100644 algorithms/dqn.py create mode 100644 algorithms/m_q_learner.py create mode 100644 algorithms/q_learner.py create mode 100644 algorithms/vdn_learner.py diff --git a/algorithms/common.py b/algorithms/common.py new file mode 100644 index 0000000..2c2f678 --- /dev/null +++ b/algorithms/common.py @@ -0,0 +1,102 @@ +from typing import NamedTuple, Union +from collections import deque, OrderedDict +import numpy as np +import random +import torch +import torch.nn as nn + + +class BaseLearner: + def __init__(self, env, n_agents, lr): + self.env = env + self.n_agents = n_agents + self.lr = lr + self.device = 'cpu' + + def to(self, device): + self.device = device + for attr, value in self.__dict__.items(): + if isinstance(value, nn.Module): + value = value.to(self.device) + return self + + +class Experience(NamedTuple): + # can be use for a single (s_t, a, r s_{t+1}) tuple + # or for a batch of tuples + observation: np.ndarray + next_observation: np.ndarray + action: np.ndarray + reward: Union[float, np.ndarray] + done : Union[bool, np.ndarray] + + +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, cer=4): + sample = random.choices(self.experience, k=k-cer) + for i in range(cer): sample += [self.experience[-i]] + observations = torch.stack([torch.from_numpy(e.observation) for e in sample], 0).float() + next_observations = torch.stack([torch.from_numpy(e.next_observation) for e in sample], 0).float() + actions = torch.tensor([e.action for e in sample]).long() + rewards = torch.tensor([e.reward for e in sample]).float().view(-1, 1) + dones = torch.tensor([e.done for e in sample]).float().view(-1, 1) + return Experience(observations, next_observations, actions, rewards, dones) + + +def soft_update(local_model, target_model, tau): + # taken from https://github.com/BY571/Munchausen-RL/blob/master/M-DQN.ipynb + for target_param, local_param in zip(target_model.parameters(), local_model.parameters()): + target_param.data.copy_(tau*local_param.data + (1.-tau)*target_param.data) + + +def mlp_maker(dims, flatten=False, activation='elu', activation_last='identity'): + activations = {'elu': nn.ELU, 'relu': nn.ReLU, + 'leaky_relu': nn.LeakyReLU, 'tanh': nn.Tanh, + 'gelu': nn.GELU, 'identity': nn.Identity} + layers = [('Flatten', nn.Flatten())] if flatten else [] + for i in range(1, len(dims)): + layers.append((f'Layer #{i - 1}: Linear', nn.Linear(dims[i - 1], dims[i]))) + activation_str = activation if i != len(dims)-1 else activation_last + layers.append((f'Layer #{i - 1}: {activation_str.capitalize()}', activations[activation_str]())) + return nn.Sequential(OrderedDict(layers)) + + +class BaseDQN(nn.Module): + def __init__(self, dims=[3*5*5, 64, 64, 9]): + super(BaseDQN, self).__init__() + self.net = mlp_maker(dims, flatten=True) + + @torch.no_grad() + def act(self, x) -> np.ndarray: + action = self.forward(x).max(-1)[1].numpy() + return action + + def forward(self, x): + return self.net(x) + + +class BaseDDQN(BaseDQN): + def __init__(self, + backbone_dims=[3*5*5, 64, 64], + value_dims=[64, 1], + advantage_dims=[64, 9]): + super(BaseDDQN, self).__init__(backbone_dims) + self.net = mlp_maker(backbone_dims, flatten=True) + self.value_head = mlp_maker(value_dims) + self.advantage_head = mlp_maker(advantage_dims) + + def forward(self, x): + features = self.net(x) + advantages = self.advantage_head(features) + values = self.value_head(features) + return values + (advantages - advantages.mean()) diff --git a/algorithms/dqn.py b/algorithms/dqn.py deleted file mode 100644 index 7ce2d20..0000000 --- a/algorithms/dqn.py +++ /dev/null @@ -1,277 +0,0 @@ -from typing import NamedTuple, Union -from collections import deque, OrderedDict -import numpy as np -import random -import torch -import torch.nn as nn -import torch.nn.functional as F - - -class Experience(NamedTuple): - # can be use for a single (s_t, a, r s_{t+1}) tuple - # or for a batch of tuples - observation: np.ndarray - next_observation: np.ndarray - action: np.ndarray - reward: Union[float, np.ndarray] - done : Union[bool, np.ndarray] - - -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, cer=4): - sample = random.choices(self.experience, k=k-cer) - for i in range(cer): sample += [self.experience[-i]] - observations = torch.stack([torch.from_numpy(e.observation) for e in sample], 0).float() - next_observations = torch.stack([torch.from_numpy(e.next_observation) for e in sample], 0).float() - actions = torch.tensor([e.action for e in sample]).long() - rewards = torch.tensor([e.reward for e in sample]).float().view(-1, 1) - dones = torch.tensor([e.done for e in sample]).float().view(-1, 1) - return Experience(observations, next_observations, actions, rewards, dones) - - - - -def soft_update(local_model, target_model, tau): - # taken from https://github.com/BY571/Munchausen-RL/blob/master/M-DQN.ipynb - for target_param, local_param in zip(target_model.parameters(), local_model.parameters()): - target_param.data.copy_(tau*local_param.data + (1.-tau)*target_param.data) - - -def mlp_maker(dims, flatten=False, activation='elu', activation_last='identity'): - activations = {'elu': nn.ELU, 'relu': nn.ReLU, - 'leaky_relu': nn.LeakyReLU, 'tanh': nn.Tanh, - 'gelu': nn.GELU, 'identity': nn.Identity} - layers = [('Flatten', nn.Flatten())] if flatten else [] - for i in range(1, len(dims)): - layers.append((f'Layer #{i - 1}: Linear', nn.Linear(dims[i - 1], dims[i]))) - activation_str = activation if i != len(dims)-1 else activation_last - layers.append((f'Layer #{i - 1}: {activation_str.capitalize()}', activations[activation_str]())) - return nn.Sequential(OrderedDict(layers)) - - -class BaseDQN(nn.Module): - def __init__(self, dims=[3*5*5, 64, 64, 9]): - super(BaseDQN, self).__init__() - self.net = mlp_maker(dims, flatten=True) - - def act(self, x) -> np.ndarray: - with torch.no_grad(): - action = self.forward(x).max(-1)[1].numpy() - return action - - def forward(self, x): - return self.net(x) - - -class BaseDDQN(BaseDQN): - def __init__(self, - backbone_dims=[3*5*5, 64, 64], - value_dims=[64,1], - advantage_dims=[64,9]): - super(BaseDDQN, self).__init__(backbone_dims) - self.net = mlp_maker(backbone_dims, flatten=True) - self.value_head = mlp_maker(value_dims) - self.advantage_head = mlp_maker(advantage_dims) - - def forward(self, x): - features = self.net(x) - advantages = self.advantage_head(features) - values = self.value_head(features) - return values + (advantages - advantages.mean()) - - -class BaseQlearner: - def __init__(self, q_net, target_q_net, env, buffer_size, target_update, eps_end, n_agents=1, - gamma=0.99, train_every_n_steps=4, n_grad_steps=1, tau=1.0, max_grad_norm=10, - 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() - soft_update(self.q_net, self.target_q_net, tau=1.0) - self.env = env - self.buffer = BaseBuffer(buffer_size) - self.target_update = target_update - 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.n_grad_steps = n_grad_steps - self.lr = lr - self.tau = tau - self.reg_weight = reg_weight - self.n_agents = n_agents - self.device = 'cpu' - self.optimizer = torch.optim.AdamW(self.q_net.parameters(), lr=self.lr) - self.max_grad_norm = max_grad_norm - self.running_reward = deque(maxlen=5) - self.running_loss = deque(maxlen=5) - self._n_updates = 0 - - def to(self, device): - self.device = device - return self - - @staticmethod - def weights_init(module, activation='leaky_relu'): - if isinstance(module, (nn.Linear, nn.Conv2d)): - nn.init.orthogonal_(module.weight, gain=torch.nn.init.calculate_gain(activation)) - if module.bias is not None: - module.bias.data.fill_(0.0) - - def anneal_eps(self, step, n_steps): - fraction = min(float(step) / int(self.exploration_fraction*n_steps), 1.0) - self.eps = 1 + fraction * (self.eps_end - 1) - - def get_action(self, obs) -> Union[int, np.ndarray]: - o = torch.from_numpy(obs).unsqueeze(0) if self.n_agents <= 1 else torch.from_numpy(obs) - if np.random.rand() > self.eps: - action = self.q_net.act(o.float()) - else: - action = np.array([self.env.action_space.sample() for _ in range(self.n_agents)]) - return action - - def learn(self, n_steps): - step = 0 - while step < n_steps: - obs, done = self.env.reset(), False - total_reward = 0 - while not done: - - action = self.get_action(obs) - - next_obs, reward, done, info = self.env.step(action if not len(action) == 1 else action[0]) - - experience = Experience(observation=obs, next_observation=next_obs, action=action, reward=reward, done=done) # do we really need to copy? - self.buffer.add(experience) - # end of step routine - obs = next_obs - step += 1 - total_reward += reward - self.anneal_eps(step, n_steps) - - if step % self.train_every_n_steps == 0: - self.train() - self._n_updates += 1 - if step % self.target_update == 0: - print('UPDATE') - soft_update(self.q_net, self.target_q_net, tau=self.tau) - - self.running_reward.append(total_reward) - if step % 10 == 0: - print(f'Step: {step} ({(step/n_steps)*100:.2f}%)\tRunning reward: {sum(list(self.running_reward))/len(self.running_reward):.2f}\t' - f' eps: {self.eps:.4f}\tRunning loss: {sum(list(self.running_loss))/len(self.running_loss):.4f}\tUpdates:{self._n_updates}') - - def _training_routine(self, obs, next_obs, action): - current_q_values = self.q_net(obs) - current_q_values = torch.gather(current_q_values, dim=-1, index=action) - next_q_values_raw = self.target_q_net(next_obs).max(dim=-1)[0].reshape(-1, 1).detach() - return current_q_values, next_q_values_raw - - def _backprop_loss(self, loss): - # log loss - self.running_loss.append(loss.item()) - # Optimize the model - self.optimizer.zero_grad() - loss.backward() - torch.nn.utils.clip_grad_norm_(self.q_net.parameters(), self.max_grad_norm) - self.optimizer.step() - - def train(self): - if len(self.buffer) < self.batch_size: return - for _ in range(self.n_grad_steps): - - experience = self.buffer.sample(self.batch_size, cer=self.train_every_n_steps) - if self.n_agents <= 1: - pred_q, target_q_raw = self._training_routine(experience.observation, - experience.next_observation, - experience.action) - - else: - pred_q, target_q_raw = torch.zeros((self.batch_size, 1)), torch.zeros((self.batch_size, 1)) - for agent_i in range(self.n_agents): - q_values, next_q_values_raw = self._training_routine(experience.observation[:, agent_i], - experience.next_observation[:, agent_i], - experience.action[:, agent_i].unsqueeze(-1)) - pred_q += q_values - target_q_raw += next_q_values_raw - target_q = experience.reward + (1 - experience.done) * self.gamma * target_q_raw - loss = torch.mean(self.reg_weight * pred_q + torch.pow(pred_q - target_q, 2)) - self._backprop_loss(loss) - - -class MunchhausenQLearner(BaseQlearner): - def __init__(self, *args, temperature=0.03, alpha=0.9, clip_l0=-1.0, **kwargs): - super(MunchhausenQLearner, self).__init__(*args, **kwargs) - assert self.n_agents == 1, 'M-DQN currently only supports single agent training' - self.temperature = temperature - self.alpha = alpha - self.clip0 = clip_l0 - - def tau_ln_pi(self, qs): - # computes log(softmax(qs/temperature)) - # Custom log-sum-exp trick from page 18 to compute the log-policy terms - v_k = qs.max(-1)[0].unsqueeze(-1) - advantage = qs - v_k - logsum = torch.logsumexp(advantage / self.temperature, -1).unsqueeze(-1) - tau_ln_pi = advantage - self.temperature * logsum - return tau_ln_pi - - def train(self): - if len(self.buffer) < self.batch_size: return - for _ in range(self.n_grad_steps): - - experience = self.buffer.sample(self.batch_size, cer=self.train_every_n_steps) - - q_target_next = self.target_q_net(experience.next_observation).detach() - tau_log_pi_next = self.tau_ln_pi(q_target_next) - - q_k_targets = self.target_q_net(experience.observation).detach() - log_pi = self.tau_ln_pi(q_k_targets) - - pi_target = F.softmax(q_target_next / self.temperature, dim=-1) - q_target = (self.gamma * (pi_target * (q_target_next - tau_log_pi_next) * (1 - experience.done)).sum(-1)).unsqueeze(-1) - - munchausen_addon = log_pi.gather(-1, experience.action) - - munchausen_reward = (experience.reward + self.alpha * torch.clamp(munchausen_addon, min=self.clip0, max=0)) - - # Compute Q targets for current states - m_q_target = munchausen_reward + q_target - - # Get expected Q values from local model - q_k = self.q_net(experience.observation) - pred_q = q_k.gather(-1, experience.action) - - # Compute loss - loss = torch.mean(self.reg_weight * pred_q + torch.pow(pred_q - m_q_target, 2)) - self._backprop_loss(loss) - - -if __name__ == '__main__': - from environments.factory.simple_factory import SimpleFactory, DirtProperties, MovementProperties - - N_AGENTS = 1 - - 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=N_AGENTS, pomdp_radius=2, max_steps=400, omit_agent_slice_in_obs=False, combin_agent_slices_in_obs=True) - - dqn, target_dqn = BaseDDQN(), BaseDDQN() - learner = MunchhausenQLearner(dqn, target_dqn, env, 40000, target_update=3500, lr=0.0007, gamma=0.99, n_agents=N_AGENTS, tau=0.95, max_grad_norm=10, - train_every_n_steps=4, eps_end=0.025, n_grad_steps=1, reg_weight=0.1, exploration_fraction=0.25, batch_size=64) - learner.learn(100000) diff --git a/algorithms/m_q_learner.py b/algorithms/m_q_learner.py new file mode 100644 index 0000000..402c68c --- /dev/null +++ b/algorithms/m_q_learner.py @@ -0,0 +1,53 @@ +import torch +import torch.nn.functional as F +from algorithms.q_learner import QLearner + + +class MQLearner(QLearner): + # Munchhausen Q-Learning + def __init__(self, *args, temperature=0.03, alpha=0.9, clip_l0=-1.0, **kwargs): + super(MQLearner, self).__init__(*args, **kwargs) + assert self.n_agents == 1, 'M-DQN currently only supports single agent training' + self.temperature = temperature + self.alpha = alpha + self.clip0 = clip_l0 + + def tau_ln_pi(self, qs): + # computes log(softmax(qs/temperature)) + # Custom log-sum-exp trick from page 18 to compute the log-policy terms + v_k = qs.max(-1)[0].unsqueeze(-1) + advantage = qs - v_k + logsum = torch.logsumexp(advantage / self.temperature, -1).unsqueeze(-1) + tau_ln_pi = advantage - self.temperature * logsum + return tau_ln_pi + + def train(self): + if len(self.buffer) < self.batch_size: return + for _ in range(self.n_grad_steps): + + experience = self.buffer.sample(self.batch_size, cer=self.train_every_n_steps) + + with torch.no_grad(): + q_target_next = self.target_q_net(experience.next_observation) + tau_log_pi_next = self.tau_ln_pi(q_target_next) + + q_k_targets = self.target_q_net(experience.observation) + log_pi = self.tau_ln_pi(q_k_targets) + + pi_target = F.softmax(q_target_next / self.temperature, dim=-1) + q_target = (self.gamma * (pi_target * (q_target_next - tau_log_pi_next) * (1 - experience.done)).sum(-1)).unsqueeze(-1) + + munchausen_addon = log_pi.gather(-1, experience.action) + + munchausen_reward = (experience.reward + self.alpha * torch.clamp(munchausen_addon, min=self.clip0, max=0)) + + # Compute Q targets for current states + m_q_target = munchausen_reward + q_target + + # Get expected Q values from local model + q_k = self.q_net(experience.observation) + pred_q = q_k.gather(-1, experience.action) + + # Compute loss + loss = torch.mean(self.reg_weight * pred_q + torch.pow(pred_q - m_q_target, 2)) + self._backprop_loss(loss) \ No newline at end of file diff --git a/algorithms/q_learner.py b/algorithms/q_learner.py new file mode 100644 index 0000000..d6ee864 --- /dev/null +++ b/algorithms/q_learner.py @@ -0,0 +1,144 @@ +from typing import Union +import gym +import torch +import torch.nn as nn +import numpy as np +from collections import deque +from pathlib import Path +import yaml +from algorithms.common import BaseLearner, BaseBuffer, soft_update, Experience + + +class QLearner(BaseLearner): + def __init__(self, q_net, target_q_net, env, buffer_size=1e5, target_update=3000, eps_end=0.05, n_agents=1, + gamma=0.99, train_every_n_steps=4, n_grad_steps=1, tau=1.0, max_grad_norm=10, weight_decay=1e-2, + exploration_fraction=0.2, batch_size=64, lr=1e-4, reg_weight=0.0, eps_start=1): + super(QLearner, self).__init__(env, n_agents, lr) + self.q_net = q_net + self.target_q_net = target_q_net + self.target_q_net.eval() + soft_update(self.q_net, self.target_q_net, tau=1.0) + self.buffer = BaseBuffer(buffer_size) + self.target_update = target_update + self.eps = eps_start + self.eps_start = eps_start + 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.n_grad_steps = n_grad_steps + self.tau = tau + self.reg_weight = reg_weight + self.weight_decay = weight_decay + self.optimizer = torch.optim.AdamW(self.q_net.parameters(), + lr=self.lr, + weight_decay=self.weight_decay) + self.max_grad_norm = max_grad_norm + self.running_reward = deque(maxlen=5) + self.running_loss = deque(maxlen=5) + self.n_updates = 0 + + def save(self, path): + path = Path(path) # no-op if already instance of Path + path.mkdir(parents=True, exist_ok=True) + hparams = {k: v for k, v in self.__dict__.items() if not(isinstance(v, BaseBuffer) or + isinstance(v, torch.optim.Optimizer) or + isinstance(v, gym.Env) or + isinstance(v, nn.Module)) + } + hparams.update({'class': self.__class__.__name__}) + with (path / 'hparams.yaml').open('w') as outfile: + yaml.dump(hparams, outfile) + torch.save(self.q_net, path / 'q_net.pt') + + def anneal_eps(self, step, n_steps): + fraction = min(float(step) / int(self.exploration_fraction*n_steps), 1.0) + self.eps = 1 + fraction * (self.eps_end - 1) + + def get_action(self, obs) -> Union[int, np.ndarray]: + o = torch.from_numpy(obs).unsqueeze(0) if self.n_agents <= 1 else torch.from_numpy(obs) + if np.random.rand() > self.eps: + action = self.q_net.act(o.float()) + else: + action = np.array([self.env.action_space.sample() for _ in range(self.n_agents)]) + return action + + def learn(self, n_steps): + step = 0 + while step < n_steps: + obs, done = self.env.reset(), False + total_reward = 0 + while not done: + + action = self.get_action(obs) + + next_obs, reward, done, info = self.env.step(action if not len(action) == 1 else action[0]) + + experience = Experience(observation=obs, next_observation=next_obs, action=action, reward=reward, done=done) # do we really need to copy? + self.buffer.add(experience) + # end of step routine + obs = next_obs + step += 1 + total_reward += reward + self.anneal_eps(step, n_steps) + + if step % self.train_every_n_steps == 0: + self.train() + self.n_updates += 1 + if step % self.target_update == 0: + print('UPDATE') + soft_update(self.q_net, self.target_q_net, tau=self.tau) + + self.running_reward.append(total_reward) + if step % 10 == 0: + print(f'Step: {step} ({(step/n_steps)*100:.2f}%)\tRunning reward: {sum(list(self.running_reward))/len(self.running_reward):.2f}\t' + f' eps: {self.eps:.4f}\tRunning loss: {sum(list(self.running_loss))/len(self.running_loss):.4f}\tUpdates:{self.n_updates}') + + def _training_routine(self, obs, next_obs, action): + current_q_values = self.q_net(obs) + current_q_values = torch.gather(current_q_values, dim=-1, index=action) + next_q_values_raw = self.target_q_net(next_obs).max(dim=-1)[0].reshape(-1, 1).detach() + return current_q_values, next_q_values_raw + + def _backprop_loss(self, loss): + # log loss + self.running_loss.append(loss.item()) + # Optimize the model + self.optimizer.zero_grad() + loss.backward() + torch.nn.utils.clip_grad_norm_(self.q_net.parameters(), self.max_grad_norm) + self.optimizer.step() + + def train(self): + if len(self.buffer) < self.batch_size: return + for _ in range(self.n_grad_steps): + experience = self.buffer.sample(self.batch_size, cer=self.train_every_n_steps) + pred_q, target_q_raw = self._training_routine(experience.observation, + experience.next_observation, + experience.action) + target_q = experience.reward + (1 - experience.done) * self.gamma * target_q_raw + loss = torch.mean(self.reg_weight * pred_q + torch.pow(pred_q - target_q, 2)) + self._backprop_loss(loss) + + + +if __name__ == '__main__': + from environments.factory.simple_factory import SimpleFactory, DirtProperties, MovementProperties + from algorithms.common import BaseDDQN + from algorithms.vdn_learner import VDNLearner + + N_AGENTS = 1 + + 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=N_AGENTS, pomdp_radius=2, max_steps=400, omit_agent_slice_in_obs=False, combin_agent_slices_in_obs=True) + + dqn, target_dqn = BaseDDQN(), BaseDDQN() + learner = QLearner(dqn, target_dqn, env, 40000, target_update=3500, lr=0.0007, gamma=0.99, n_agents=N_AGENTS, tau=0.95, max_grad_norm=10, + train_every_n_steps=4, eps_end=0.025, n_grad_steps=1, reg_weight=0.1, exploration_fraction=0.25, batch_size=64) + #learner.save(Path(__file__).parent / 'test' / 'testexperiment1337') + learner.learn(100000) diff --git a/algorithms/vdn_learner.py b/algorithms/vdn_learner.py new file mode 100644 index 0000000..f50c6ab --- /dev/null +++ b/algorithms/vdn_learner.py @@ -0,0 +1,23 @@ +import torch +from algorithms.q_learner import QLearner + + +class VDNLearner(QLearner): + def __init__(self, *args, **kwargs): + super(VDNLearner, self).__init__(*args, **kwargs) + assert self.n_agents >= 2, 'VDN requires more than one agent, use QLearner instead' + + def train(self): + if len(self.buffer) < self.batch_size: return + for _ in range(self.n_grad_steps): + experience = self.buffer.sample(self.batch_size, cer=self.train_every_n_steps) + pred_q, target_q_raw = torch.zeros((self.batch_size, 1)), torch.zeros((self.batch_size, 1)) + for agent_i in range(self.n_agents): + q_values, next_q_values_raw = self._training_routine(experience.observation[:, agent_i], + experience.next_observation[:, agent_i], + experience.action[:, agent_i].unsqueeze(-1)) + pred_q += q_values + target_q_raw += next_q_values_raw + target_q = experience.reward + (1 - experience.done) * self.gamma * target_q_raw + loss = torch.mean(self.reg_weight * pred_q + torch.pow(pred_q - target_q, 2)) + self._backprop_loss(loss) \ No newline at end of file From 456e48f2e0e48cdeda254127df91975781afe02b Mon Sep 17 00:00:00 2001 From: romue Date: Fri, 25 Jun 2021 15:42:55 +0200 Subject: [PATCH 8/9] added individual eps-greedy for VDN --- algorithms/common.py | 18 ++++++++++++++ algorithms/q_learner.py | 4 ++-- algorithms/qtran_learner.py | 48 +++++++++++++++++++++++++++++++++++++ algorithms/vdn_learner.py | 17 +++++++++++++ 4 files changed, 85 insertions(+), 2 deletions(-) create mode 100644 algorithms/qtran_learner.py diff --git a/algorithms/common.py b/algorithms/common.py index 2c2f678..ddc1136 100644 --- a/algorithms/common.py +++ b/algorithms/common.py @@ -71,6 +71,7 @@ def mlp_maker(dims, flatten=False, activation='elu', activation_last='identity') return nn.Sequential(OrderedDict(layers)) + class BaseDQN(nn.Module): def __init__(self, dims=[3*5*5, 64, 64, 9]): super(BaseDQN, self).__init__() @@ -100,3 +101,20 @@ class BaseDDQN(BaseDQN): advantages = self.advantage_head(features) values = self.value_head(features) return values + (advantages - advantages.mean()) + + +class QTRANtestNet(nn.Module): + def __init__(self, backbone_dims=[3*5*5, 64, 64], q_head=[64, 9]): + super(QTRANtestNet, self).__init__() + self.backbone = mlp_maker(backbone_dims, flatten=True, activation_last='elu') + self.q_head = mlp_maker(q_head) + + def forward(self, x): + features = self.backbone(x) + qs = self.q_head(features) + return qs, features + + @torch.no_grad() + def act(self, x) -> np.ndarray: + action = self.forward(x)[0].max(-1)[1].numpy() + return action \ No newline at end of file diff --git a/algorithms/q_learner.py b/algorithms/q_learner.py index d6ee864..10d34a2 100644 --- a/algorithms/q_learner.py +++ b/algorithms/q_learner.py @@ -128,7 +128,7 @@ if __name__ == '__main__': from algorithms.common import BaseDDQN from algorithms.vdn_learner import VDNLearner - N_AGENTS = 1 + N_AGENTS = 2 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) @@ -138,7 +138,7 @@ if __name__ == '__main__': env = SimpleFactory(dirt_properties=dirt_props, movement_properties=move_props, n_agents=N_AGENTS, pomdp_radius=2, max_steps=400, omit_agent_slice_in_obs=False, combin_agent_slices_in_obs=True) dqn, target_dqn = BaseDDQN(), BaseDDQN() - learner = QLearner(dqn, target_dqn, env, 40000, target_update=3500, lr=0.0007, gamma=0.99, n_agents=N_AGENTS, tau=0.95, max_grad_norm=10, + learner = VDNLearner(dqn, target_dqn, env, 40000, target_update=3500, lr=0.0007, gamma=0.99, n_agents=N_AGENTS, tau=0.95, max_grad_norm=10, train_every_n_steps=4, eps_end=0.025, n_grad_steps=1, reg_weight=0.1, exploration_fraction=0.25, batch_size=64) #learner.save(Path(__file__).parent / 'test' / 'testexperiment1337') learner.learn(100000) diff --git a/algorithms/qtran_learner.py b/algorithms/qtran_learner.py new file mode 100644 index 0000000..fc6cc24 --- /dev/null +++ b/algorithms/qtran_learner.py @@ -0,0 +1,48 @@ +import torch +from algorithms.q_learner import QLearner + + +class QTRANLearner(QLearner): + def __init__(self, *args, weight_opt=1., weigt_nopt=1., **kwargs): + super(QTRANLearner, self).__init__(*args, **kwargs) + assert self.n_agents >= 2, 'QTRANLearner requires more than one agent, use QLearner instead' + self.weight_opt = weight_opt + self.weigt_nopt = weigt_nopt + + def _training_routine(self, obs, next_obs, action): + # todo remove - is inherited - only used while implementing qtran + current_q_values = self.q_net(obs) + current_q_values = torch.gather(current_q_values, dim=-1, index=action) + next_q_values_raw = self.target_q_net(next_obs).max(dim=-1)[0].reshape(-1, 1).detach() + return current_q_values, next_q_values_raw + + def local_qs(self, observations, actions): + Q_jt = torch.zeros_like(actions) # placeholder to sum up individual q values + features = [] + for agent_i in range(self.n_agents): + q_values_agent_i, features_agent_i = self.q_net(observations[:, agent_i]) # Individual action-value network + q_values_agent_i = torch.gather(q_values_agent_i, dim=-1, index=actions[:, agent_i].unsqueeze(-1)) + Q_jt += q_values_agent_i + features.append(features_agent_i) + feature_sum = torch.stack(features, 0).sum(0) # (n_agents x hdim) -> hdim + return Q_jt + + def train(self): + if len(self.buffer) < self.batch_size: return + for _ in range(self.n_grad_steps): + experience = self.buffer.sample(self.batch_size, cer=self.train_every_n_steps) + + Q_jt_prime = self.local_qs(experience.observation, experience.action) # sum of individual q-vals + Q_jt = None + V_jt = None + + pred_q, target_q_raw = torch.zeros((self.batch_size, 1)), torch.zeros((self.batch_size, 1)) + for agent_i in range(self.n_agents): + q_values, next_q_values_raw = self._training_routine(experience.observation[:, agent_i], + experience.next_observation[:, agent_i], + experience.action[:, agent_i].unsqueeze(-1)) + pred_q += q_values + target_q_raw += next_q_values_raw + target_q = experience.reward + (1 - experience.done) * self.gamma * target_q_raw + loss = torch.mean(self.reg_weight * pred_q + torch.pow(pred_q - target_q, 2)) + self._backprop_loss(loss) \ No newline at end of file diff --git a/algorithms/vdn_learner.py b/algorithms/vdn_learner.py index f50c6ab..504adb0 100644 --- a/algorithms/vdn_learner.py +++ b/algorithms/vdn_learner.py @@ -1,4 +1,6 @@ +from typing import Union import torch +import numpy as np from algorithms.q_learner import QLearner @@ -7,6 +9,21 @@ class VDNLearner(QLearner): super(VDNLearner, self).__init__(*args, **kwargs) assert self.n_agents >= 2, 'VDN requires more than one agent, use QLearner instead' + def get_action(self, obs) -> Union[int, np.ndarray]: + o = torch.from_numpy(obs).unsqueeze(0) if self.n_agents <= 1 else torch.from_numpy(obs) + eps = np.random.rand(self.n_agents) + greedy = eps > self.eps + agent_actions = None + actions = [] + for i in range(self.n_agents): + if greedy[i]: + if agent_actions is None: agent_actions = self.q_net.act(o.float()) + action = agent_actions[i] + else: + action = self.env.action_space.sample() + actions.append(action) + return np.array(actions) + def train(self): if len(self.buffer) < self.batch_size: return for _ in range(self.n_grad_steps): From 87f762c78c8600c508a33260c07afbc5303eebd6 Mon Sep 17 00:00:00 2001 From: romue Date: Tue, 29 Jun 2021 16:40:30 +0200 Subject: [PATCH 9/9] added individual eps-greedy for VDN --- algorithms/common.py | 96 ++++++++++++++++---- algorithms/q_learner.py | 52 ++++------- algorithms/udr_learner.py | 178 ++++++++++++++++++++++++++++++++++++++ 3 files changed, 272 insertions(+), 54 deletions(-) create mode 100644 algorithms/udr_learner.py diff --git a/algorithms/common.py b/algorithms/common.py index ddc1136..1749f7e 100644 --- a/algorithms/common.py +++ b/algorithms/common.py @@ -6,21 +6,6 @@ import torch import torch.nn as nn -class BaseLearner: - def __init__(self, env, n_agents, lr): - self.env = env - self.n_agents = n_agents - self.lr = lr - self.device = 'cpu' - - def to(self, device): - self.device = device - for attr, value in self.__dict__.items(): - if isinstance(value, nn.Module): - value = value.to(self.device) - return self - - class Experience(NamedTuple): # can be use for a single (s_t, a, r s_{t+1}) tuple # or for a batch of tuples @@ -29,6 +14,84 @@ class Experience(NamedTuple): action: np.ndarray reward: Union[float, np.ndarray] done : Union[bool, np.ndarray] + episode: int = -1 + + +class BaseLearner: + def __init__(self, env, n_agents=1, train_every=('step', 4), n_grad_steps=1): + assert train_every[0] in ['step', 'episode'], 'train_every[0] must be one of ["step", "episode"]' + self.env = env + self.n_agents = n_agents + self.n_grad_steps = n_grad_steps + self.train_every = train_every + self.device = 'cpu' + self.n_updates = 0 + self.step = 0 + self.episode_step = 0 + self.episode = 0 + self.running_reward = deque(maxlen=5) + + def to(self, device): + self.device = device + for attr, value in self.__dict__.items(): + if isinstance(value, nn.Module): + value = value.to(self.device) + return self + + def get_action(self, obs) -> Union[int, np.ndarray]: + pass + + def on_new_experience(self, experience): + pass + + def on_step_end(self, n_steps): + pass + + def on_episode_end(self, n_steps): + pass + + def train(self): + pass + + def learn(self, n_steps): + train_type, train_freq = self.train_every + while self.step < n_steps: + obs, done = self.env.reset(), False + total_reward = 0 + self.episode_step = 0 + while not done: + + action = self.get_action(obs) + + next_obs, reward, done, info = self.env.step(action if not len(action) == 1 else action[0]) + + experience = Experience(observation=obs, next_observation=next_obs, + action=action, reward=reward, + done=done, episode=self.episode) # do we really need to copy? + self.on_new_experience(experience) + # end of step routine + obs = next_obs + total_reward += reward + self.step += 1 + self.episode_step += 1 + self.on_step_end(n_steps) + if train_type == 'step' and (self.step % train_freq == 0): + self.train() + self.n_updates += 1 + self.on_episode_end(n_steps) + if train_type == 'episode' and (self.episode % train_freq == 0): + self.train() + self.n_updates += 1 + + self.running_reward.append(total_reward) + self.episode += 1 + try: + if self.step % 10 == 0: + print( + f'Step: {self.step} ({(self.step / n_steps) * 100:.2f}%)\tRunning reward: {sum(list(self.running_reward)) / len(self.running_reward):.2f}\t' + f' eps: {self.eps:.4f}\tRunning loss: {sum(list(self.running_loss)) / len(self.running_loss):.4f}\tUpdates:{self.n_updates}') + except Exception as e: + pass class BaseBuffer: @@ -60,7 +123,7 @@ def soft_update(local_model, target_model, tau): def mlp_maker(dims, flatten=False, activation='elu', activation_last='identity'): - activations = {'elu': nn.ELU, 'relu': nn.ReLU, + activations = {'elu': nn.ELU, 'relu': nn.ReLU, 'sigmoid': nn.Sigmoid, 'leaky_relu': nn.LeakyReLU, 'tanh': nn.Tanh, 'gelu': nn.GELU, 'identity': nn.Identity} layers = [('Flatten', nn.Flatten())] if flatten else [] @@ -71,7 +134,6 @@ def mlp_maker(dims, flatten=False, activation='elu', activation_last='identity') return nn.Sequential(OrderedDict(layers)) - class BaseDQN(nn.Module): def __init__(self, dims=[3*5*5, 64, 64, 9]): super(BaseDQN, self).__init__() diff --git a/algorithms/q_learner.py b/algorithms/q_learner.py index 10d34a2..06a3384 100644 --- a/algorithms/q_learner.py +++ b/algorithms/q_learner.py @@ -11,9 +11,9 @@ from algorithms.common import BaseLearner, BaseBuffer, soft_update, Experience class QLearner(BaseLearner): def __init__(self, q_net, target_q_net, env, buffer_size=1e5, target_update=3000, eps_end=0.05, n_agents=1, - gamma=0.99, train_every_n_steps=4, n_grad_steps=1, tau=1.0, max_grad_norm=10, weight_decay=1e-2, + gamma=0.99, train_every=('step', 4), n_grad_steps=1, tau=1.0, max_grad_norm=10, weight_decay=1e-2, exploration_fraction=0.2, batch_size=64, lr=1e-4, reg_weight=0.0, eps_start=1): - super(QLearner, self).__init__(env, n_agents, lr) + super(QLearner, self).__init__(env, n_agents, train_every, n_grad_steps) self.q_net = q_net self.target_q_net = target_q_net self.target_q_net.eval() @@ -26,11 +26,10 @@ class QLearner(BaseLearner): self.exploration_fraction = exploration_fraction self.batch_size = batch_size self.gamma = gamma - self.train_every_n_steps = train_every_n_steps - self.n_grad_steps = n_grad_steps self.tau = tau self.reg_weight = reg_weight self.weight_decay = weight_decay + self.lr = lr self.optimizer = torch.optim.AdamW(self.q_net.parameters(), lr=self.lr, weight_decay=self.weight_decay) @@ -64,36 +63,14 @@ class QLearner(BaseLearner): action = np.array([self.env.action_space.sample() for _ in range(self.n_agents)]) return action - def learn(self, n_steps): - step = 0 - while step < n_steps: - obs, done = self.env.reset(), False - total_reward = 0 - while not done: + def on_new_experience(self, experience): + self.buffer.add(experience) - action = self.get_action(obs) - - next_obs, reward, done, info = self.env.step(action if not len(action) == 1 else action[0]) - - experience = Experience(observation=obs, next_observation=next_obs, action=action, reward=reward, done=done) # do we really need to copy? - self.buffer.add(experience) - # end of step routine - obs = next_obs - step += 1 - total_reward += reward - self.anneal_eps(step, n_steps) - - if step % self.train_every_n_steps == 0: - self.train() - self.n_updates += 1 - if step % self.target_update == 0: - print('UPDATE') - soft_update(self.q_net, self.target_q_net, tau=self.tau) - - self.running_reward.append(total_reward) - if step % 10 == 0: - print(f'Step: {step} ({(step/n_steps)*100:.2f}%)\tRunning reward: {sum(list(self.running_reward))/len(self.running_reward):.2f}\t' - f' eps: {self.eps:.4f}\tRunning loss: {sum(list(self.running_loss))/len(self.running_loss):.4f}\tUpdates:{self.n_updates}') + def on_step_end(self, n_steps): + self.anneal_eps(self.step, n_steps) + if self.step % self.target_update == 0: + print('UPDATE') + soft_update(self.q_net, self.target_q_net, tau=self.tau) def _training_routine(self, obs, next_obs, action): current_q_values = self.q_net(obs) @@ -113,7 +90,7 @@ class QLearner(BaseLearner): def train(self): if len(self.buffer) < self.batch_size: return for _ in range(self.n_grad_steps): - experience = self.buffer.sample(self.batch_size, cer=self.train_every_n_steps) + experience = self.buffer.sample(self.batch_size, cer=self.train_every[-1]) pred_q, target_q_raw = self._training_routine(experience.observation, experience.next_observation, experience.action) @@ -127,8 +104,9 @@ if __name__ == '__main__': from environments.factory.simple_factory import SimpleFactory, DirtProperties, MovementProperties from algorithms.common import BaseDDQN from algorithms.vdn_learner import VDNLearner + from algorithms.udr_learner import UDRLearner - N_AGENTS = 2 + N_AGENTS = 1 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) @@ -138,7 +116,7 @@ if __name__ == '__main__': env = SimpleFactory(dirt_properties=dirt_props, movement_properties=move_props, n_agents=N_AGENTS, pomdp_radius=2, max_steps=400, omit_agent_slice_in_obs=False, combin_agent_slices_in_obs=True) dqn, target_dqn = BaseDDQN(), BaseDDQN() - learner = VDNLearner(dqn, target_dqn, env, 40000, target_update=3500, lr=0.0007, gamma=0.99, n_agents=N_AGENTS, tau=0.95, max_grad_norm=10, - train_every_n_steps=4, eps_end=0.025, n_grad_steps=1, reg_weight=0.1, exploration_fraction=0.25, batch_size=64) + learner = QLearner(dqn, target_dqn, env, 40000, target_update=3500, lr=0.0007, gamma=0.99, n_agents=N_AGENTS, tau=0.95, max_grad_norm=10, + train_every=('step', 4), eps_end=0.025, n_grad_steps=1, reg_weight=0.1, exploration_fraction=0.25, batch_size=64) #learner.save(Path(__file__).parent / 'test' / 'testexperiment1337') learner.learn(100000) diff --git a/algorithms/udr_learner.py b/algorithms/udr_learner.py new file mode 100644 index 0000000..b99f10f --- /dev/null +++ b/algorithms/udr_learner.py @@ -0,0 +1,178 @@ +import random +from typing import Union, List +from collections import deque +import numpy as np +import torch +import torch.nn as nn +from algorithms.common import BaseBuffer, Experience, BaseLearner, BaseDQN, mlp_maker +from collections import defaultdict + + +class UDRLBuffer(BaseBuffer): + def __init__(self, size): + super(UDRLBuffer, self).__init__(0) + self.experience = defaultdict(list) + self.size = size + + def add(self, experience): + self.experience[experience.episode].append(experience) + if len(self.experience) > self.size: + self.sort_and_prune() + + def select_time_steps(self, episode: List[Experience]): + T = len(episode) # max horizon + t1 = random.randint(0, T - 1) + t2 = random.randint(t1 + 1, T) + return t1, t2, T + + def sort_and_prune(self): + scores = [] + for k, episode_experience in self.experience.items(): + r = sum([e.reward for e in episode_experience]) + scores.append((r, k)) + sorted_scores = sorted(scores, reverse=True) + return sorted_scores + + def sample(self, batch_size, cer=0): + random_episode_keys = random.choices(list(self.experience.keys()), k=batch_size) + lsts = (obs, desired_rewards, horizons, actions) = [], [], [], [] + for ek in random_episode_keys: + episode = self.experience[ek] + t1, t2, T = self.select_time_steps(episode) + t2 = T # TODO only good for episodic envs + observation = episode[t1].observation + desired_reward = sum([experience.reward for experience in episode[t1:t2]]) + horizon = t2 - t1 + action = episode[t1].action + for lst, val in zip(lsts, [observation, desired_reward, horizon, action]): + lst.append(val) + return (torch.stack([torch.from_numpy(o) for o in obs], 0).float(), + torch.tensor(desired_rewards).view(-1, 1).float(), + torch.tensor(horizons).view(-1, 1).float(), + torch.tensor(actions)) + + +class UDRLearner(BaseLearner): + # Upside Down Reinforcement Learner + def __init__(self, env, desired_reward, desired_horizon, + behavior_fn=None, buffer_size=100, n_warm_up_episodes=8, best_x=20, + batch_size=128, lr=1e-3, n_agents=1, train_every=('episode', 4), n_grad_steps=1): + super(UDRLearner, self).__init__(env, n_agents, train_every, n_grad_steps) + assert self.n_agents == 1, 'UDRL currently only supports single agent training' + self.behavior_fn = behavior_fn + self.buffer_size = buffer_size + self.n_warm_up_episodes = n_warm_up_episodes + self.buffer = UDRLBuffer(buffer_size) + self.batch_size = batch_size + self.mode = 'train' + self.best_x = best_x + self.desired_reward = desired_reward + self.desired_horizon = desired_horizon + self.lr = lr + self.optimizer = torch.optim.AdamW(self.behavior_fn.parameters(), lr=lr) + + self.running_loss = deque(maxlen=self.n_grad_steps*5) + + def sample_exploratory_commands(self): + top_x = self.buffer.sort_and_prune()[:self.best_x] + # The exploratory desired horizon dh0 is set to the mean of the lengths of the selected episodes + new_desired_horizon = np.mean([len(self.buffer.experience[k]) for _, k in top_x]) + # save all top_X cumulative returns in a list + returns = [r for r, _ in top_x] + # from these returns calc the mean and std + mean_returns = np.mean([r for r, _ in top_x]) + std_returns = np.std(returns) + # sample desired reward from a uniform distribution given the mean and the std + new_desired_reward = np.random.uniform(mean_returns, mean_returns + std_returns) + self.exploratory_commands = (new_desired_reward, new_desired_horizon) + return torch.tensor([[new_desired_reward]]).float(), torch.tensor([[new_desired_horizon]]).float() + + def on_new_experience(self, experience): + self.buffer.add(experience) + self.desired_reward = self.desired_reward - torch.tensor(experience.reward).float().view(1, 1) + + def on_step_end(self, n_steps): + one = torch.tensor([1.]).float().view(1, 1) + self.desired_horizon -= one + self.desired_horizon = self.desired_horizon if self.desired_horizon >= 1. else one + + def on_episode_end(self, n_steps): + self.desired_reward, self.desired_horizon = self.sample_exploratory_commands() + + def get_action(self, obs) -> Union[int, np.ndarray]: + o = torch.from_numpy(obs).unsqueeze(0) if self.n_agents <= 1 else torch.from_numpy(obs) + bf_out = self.behavior_fn(o.float(), self.desired_reward, self.desired_horizon) + dist = torch.distributions.Categorical(bf_out) + sample = dist.sample() + return [sample.item()]#[self.env.action_space.sample()] + + def _backprop_loss(self, loss): + # log loss + self.running_loss.append(loss.item()) + # Optimize the model + self.optimizer.zero_grad() + loss.backward() + #torch.nn.utils.clip_grad_norm_(self.behavior_fn.parameters(), 10) + self.optimizer.step() + + def train(self): + if len(self.buffer) < self.n_warm_up_episodes: return + for _ in range(self.n_grad_steps): + experience = self.buffer.sample(self.batch_size) + bf_out = self.behavior_fn(*experience[:3]) + labels = experience[-1] + #print(labels.shape) + loss = nn.CrossEntropyLoss()(bf_out, labels.squeeze()) + mean_entropy = torch.distributions.Categorical(bf_out).entropy().mean() + self._backprop_loss(loss - 0.03*mean_entropy) + print(f'Running loss: {np.mean(list(self.running_loss)):.3f}\tRunning reward: {np.mean(self.running_reward):.2f}' + f'\td_r: {self.desired_reward.item():.2f}\ttd_h: {self.desired_horizon.item()}') + + +class BF(BaseDQN): + def __init__(self, dims=[5*5*3, 64]): + super(BF, self).__init__(dims) + self.net = mlp_maker(dims, activation_last='identity') + self.command_net = mlp_maker([2, 64], activation_last='sigmoid') + self.common_branch = mlp_maker([64, 64, 64, 9]) + + + def forward(self, observation, desired_reward, horizon): + command = torch.cat((desired_reward*(0.02), horizon*(0.01)), dim=-1) + obs_out = self.net(torch.flatten(observation, start_dim=1)) + command_out = self.command_net(command) + combined = obs_out*command_out + out = self.common_branch(combined) + return torch.softmax(out, -1) + + +if __name__ == '__main__': + from environments.factory.simple_factory import SimpleFactory, DirtProperties, MovementProperties + from algorithms.common import BaseDDQN + from algorithms.vdn_learner import VDNLearner + + N_AGENTS = 1 + + 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=N_AGENTS, pomdp_radius=2, + max_steps=400, omit_agent_slice_in_obs=False, combin_agent_slices_in_obs=True) + + bf = BF() + desired_reward = torch.tensor([200.]).view(1, 1).float() + desired_horizon = torch.tensor([400.]).view(1, 1).float() + learner = UDRLearner(env, behavior_fn=bf, + train_every=('episode', 2), + buffer_size=40, + best_x=10, + lr=1e-3, + batch_size=64, + n_warm_up_episodes=12, + n_grad_steps=4, + desired_reward=desired_reward, + desired_horizon=desired_horizon) + #learner.save(Path(__file__).parent / 'test' / 'testexperiment1337') + learner.learn(500000)