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)