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)