Feat: IQN

Author: choru-k

README.md

@@ -35,7 +35,7 @@ So you can run this example in your computer(maybe it take just only 1~2 minitue
 
 ## Distributional DQN
 - [x] QRDQN [[18]](#reference)
-- [ ] IQN [[19]](#reference)
+- [x] IQN [[19]](#reference)
 
 ## Exploration
 - [ ] ICM [[22]](#refercence)

distributional/2-IQN/config.py

@@ -0,0 +1,18 @@
+import torch
+
+env_name = 'CartPole-v1'
+gamma = 0.99
+batch_size = 32
+lr = 0.001
+initial_exploration = 1000
+goal_score = 200
+log_interval = 10
+update_target = 100
+replay_memory_capacity = 1000
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+num_quantile_sample = 32
+num_tau_sample = 16
+num_tau_prime_sample = 8
+quantile_embedding_dim = 64

distributional/2-IQN/memory.py

@@ -0,0 +1,22 @@
+import random
+from collections import namedtuple, deque
+
+
+Transition = namedtuple('Transition', ('state', 'next_state', 'action', 'reward', 'mask'))
+
+
+class Memory(object):
+    def __init__(self, capacity):
+        self.memory = deque(maxlen=capacity)
+        self.capacity = capacity
+
+    def push(self, state, next_state, action, reward, mask):
+        self.memory.append(Transition(state, next_state, action, reward, mask))
+
+    def sample(self, batch_size):
+        transitions = random.sample(self.memory, batch_size)
+        batch = Transition(*zip(*transitions))
+        return batch
+
+    def __len__(self):
+        return len(self.memory)

distributional/2-IQN/model.py

@@ -0,0 +1,83 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+
+from config import batch_size, gamma, quantile_embedding_dim, num_tau_sample, num_tau_prime_sample, num_quantile_sample
+
+class QRDQN(nn.Module):
+    def __init__(self, num_inputs, num_outputs):
+        super(QRDQN, self).__init__()
+        self.num_inputs = num_inputs
+        self.num_outputs = num_outputs
+
+        self.fc1 = nn.Linear(num_inputs, 128)
+        self.fc2 = nn.Linear(128, num_outputs)
+        self.phi = nn.Linear(quantile_embedding_dim, 128)
+
+        for m in self.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.xavier_uniform(m.weight)
+
+    def forward(self, state, tau, num_quantiles):
+        input_size = state.size()[0] # batch_size(train) or 1(get_action)
+        tau = tau.expand(input_size * num_quantiles, quantile_embedding_dim)
+        pi_mtx = torch.Tensor(np.pi * np.arange(0, quantile_embedding_dim)).expand(input_size * num_quantiles, quantile_embedding_dim)
+        cos_tau = torch.cos(tau * pi_mtx)
+
+        phi = self.phi(cos_tau)
+        phi = F.relu(phi)
+
+        state_tile = state.expand(input_size, num_quantiles, self.num_inputs)
+        state_tile = state_tile.flatten().view(-1, self.num_inputs)
+        
+        x = F.relu(self.fc1(state_tile))
+        x = self.fc2(x * phi)
+        z = x.view(-1, num_quantiles, self.num_outputs)
+
+        z = z.transpose(1, 2) # [input_size, num_output, num_quantile]
+        return z
+
+    def get_action(self, state):
+        tau = torch.Tensor(np.random.rand(num_quantile_sample, 1) * 0.5) # CVaR
+        z = self.forward(state, tau, num_quantile_sample)
+        q = z.mean(dim=2, keepdim=True)
+        action = torch.argmax(q)
+        return action.item()
+
+    @classmethod
+    def train_model(cls, online_net, target_net, optimizer, batch):
+        states = torch.stack(batch.state)
+        next_states = torch.stack(batch.next_state)
+        actions = torch.Tensor(batch.action).long()
+        rewards = torch.Tensor(batch.reward)
+        masks = torch.Tensor(batch.mask)
+
+        tau = torch.Tensor(np.random.rand(batch_size * num_tau_sample, 1))
+        z = online_net(states, tau, num_tau_sample)
+        action = actions.unsqueeze(1).unsqueeze(1).expand(-1, 1, num_tau_sample)
+        z_a = z.gather(1, action).squeeze(1)
+
+        tau_prime = torch.Tensor(np.random.rand(batch_size * num_tau_prime_sample, 1))
+        next_z = target_net(next_states, tau_prime, num_tau_prime_sample)
+        next_action = next_z.mean(dim=2).max(1)[1]
+        next_action = next_action.unsqueeze(1).unsqueeze(1).expand(batch_size, 1, num_tau_prime_sample)
+        next_z_a = next_z.gather(1, next_action).squeeze(1)
+
+        T_z = rewards.unsqueeze(1) + gamma * next_z_a * masks.unsqueeze(1)
+
+        T_z_tile = T_z.view(-1, num_tau_prime_sample, 1).expand(-1, num_tau_prime_sample, num_tau_sample)
+        z_a_tile = z_a.view(-1, 1, num_tau_sample).expand(-1, num_tau_prime_sample, num_tau_sample)
+        
+        error_loss = T_z_tile - z_a_tile
+        huber_loss = nn.SmoothL1Loss(reduction='none')(T_z_tile, z_a_tile)
+        tau = torch.arange(0, 1, 1 / num_tau_sample).view(1, num_tau_sample)
+        
+        loss = (tau - (error_loss < 0).float()).abs() * huber_loss
+        loss = loss.mean(dim=2).sum(dim=1).mean()
+
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+
+        return loss

distributional/2-IQN/train.py

@@ -0,0 +1,102 @@
+import os
+import sys
+import gym
+import random
+import argparse
+import numpy as np
+
+import torch
+import torch.optim as optim
+import torch.nn.functional as F
+from tensorboardX import SummaryWriter
+
+from model import QRDQN
+from memory import Memory
+
+from config import env_name, initial_exploration, batch_size, update_target, goal_score, log_interval, device, replay_memory_capacity, lr
+
+
+def get_action(state, target_net, epsilon, env):
+    if np.random.rand() <= epsilon:
+        return env.action_space.sample()
+    else:
+        return target_net.get_action(state)
+
+def update_target_model(online_net, target_net):
+    # Target <- Net
+    target_net.load_state_dict(online_net.state_dict())
+
+
+def main():
+    env = gym.make(env_name)
+    env.seed(500)
+    torch.manual_seed(500)
+
+    num_inputs = env.observation_space.shape[0]
+    num_actions = env.action_space.n
+    print('state size:', num_inputs)
+    print('action size:', num_actions)
+
+    online_net = QRDQN(num_inputs, num_actions)
+    target_net = QRDQN(num_inputs, num_actions)
+    update_target_model(online_net, target_net)
+
+    optimizer = optim.Adam(online_net.parameters(), lr=lr)
+    writer = SummaryWriter('logs')
+
+    online_net.to(device)
+    target_net.to(device)
+    online_net.train()
+    target_net.train()
+    memory = Memory(replay_memory_capacity)
+    running_score = 0
+    epsilon = 1.0
+    steps = 0
+    loss = 0
+
+    for e in range(3000):
+        done = False
+
+        score = 0
+        state = env.reset()
+        state = torch.Tensor(state)
+        state = state.unsqueeze(0)
+
+        while not done:
+            steps += 1
+            action = get_action(state, target_net, epsilon, env)
+            next_state, reward, done, _ = env.step(action)
+
+            next_state = torch.Tensor(next_state)
+            next_state = next_state.unsqueeze(0)
+
+            mask = 0 if done else 1
+            reward = reward if not done or score == 499 else -1
+            memory.push(state, next_state, action, reward, mask)
+
+            score += reward
+            state = next_state
+
+            if steps > initial_exploration:
+                epsilon -= 0.00005
+                epsilon = max(epsilon, 0.1)
+
+                batch = memory.sample(batch_size)
+                loss = QRDQN.train_model(online_net, target_net, optimizer, batch)
+
+                if steps % update_target == 0:
+                    update_target_model(online_net, target_net)
+
+        score = score if score == 500.0 else score + 1
+        running_score = 0.99 * running_score + 0.01 * score
+        if e % log_interval == 0:
+            print('{} episode | score: {:.2f} | epsilon: {:.2f}'.format(
+                e, running_score, epsilon))
+            writer.add_scalar('log/score', float(running_score), e)
+            writer.add_scalar('log/loss', float(loss), e)
+
+        if running_score > goal_score:
+            break
+
+if __name__=="__main__":
+    main()