diff --git a/.gitignore b/.gitignore
index 5a14782..d9eb23a 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,3 +1,4 @@
 train_data/*
 __pycache__
 
+.idea
diff --git a/agent/__init__.py b/agents/__init__.py
similarity index 100%
rename from agent/__init__.py
rename to agents/__init__.py
diff --git a/agent/alphago.py b/agents/alphago.py
similarity index 96%
rename from agent/alphago.py
rename to agents/alphago.py
index 4def874..167f5ac 100644
--- a/agent/alphago.py
+++ b/agents/alphago.py
@@ -1,16 +1,6 @@
-# tag::alphago_imports[]
 import numpy as np
-from tugo.agent.base import Agent
-from tugo.goboard_fast import Move
-from tugo import kerasutil
-import operator
-# end::alphago_imports[]
-
-
-__all__ = [
-    'AlphaGoNode',
-    'AlphaGoMCTS'
-]
+from tugo.agents.base import Agent
+from game_logic.goboard_fast import Move
 
 
 # tag::init_alphago_node[]
diff --git a/agent/base.py b/agents/base.py
similarity index 74%
rename from agent/base.py
rename to agents/base.py
index 161dedd..c92f2cd 100644
--- a/agent/base.py
+++ b/agents/base.py
@@ -1,16 +1,10 @@
-__all__ = [
-    'Agent',
-]
 
-
-# tag::agent[]
 class Agent:
     def __init__(self):
         pass
 
     def select_move(self, game_state):
         raise NotImplementedError()
-# end::agent[]
 
     def diagnostics(self):
         return {}
diff --git a/agent/helpers.py b/agents/helpers.py
similarity index 91%
rename from agent/helpers.py
rename to agents/helpers.py
index 5cb52e4..2a0426f 100644
--- a/agent/helpers.py
+++ b/agents/helpers.py
@@ -1,10 +1,4 @@
-# tag::helpersimport[]
-from tugo.gotypes import Point
-# end::helpersimport[]
-
-__all__ = [
-    'is_point_an_eye',
-]
+from game_logic.gotypes import Point
 
 
 # tag::eye[]
diff --git a/agent/helpers_fast.py b/agents/helpers_fast.py
similarity index 94%
rename from agent/helpers_fast.py
rename to agents/helpers_fast.py
index 6f9f28b..0a142c0 100644
--- a/agent/helpers_fast.py
+++ b/agents/helpers_fast.py
@@ -1,8 +1,4 @@
-from tugo.gotypes import Point
-
-__all__ = [
-    'is_point_an_eye',
-]
+from game_logic.gotypes import Point
 
 
 def is_point_an_eye(board, point, color):
diff --git a/agent/naive.py b/agents/naive.py
similarity index 77%
rename from agent/naive.py
rename to agents/naive.py
index a838e45..d220adc 100644
--- a/agent/naive.py
+++ b/agents/naive.py
@@ -1,16 +1,10 @@
-# tag::randombotimports[]
 import random
-from tugo.agent.base import Agent
-from tugo.agent.helpers import is_point_an_eye
+from tugo.agents.base import Agent
+from tugo.agents.helpers import is_point_an_eye
 from tugo.goboard_slow import Move
-from tugo.gotypes import Point
-# end::randombotimports[]
+from game_logic.gotypes import Point
 
 
-__all__ = ['RandomBot']
-
-
-# tag::random_bot[]
 class RandomBot(Agent):
     def select_move(self, game_state):
         """Choose a random valid move that preserves our own eyes."""
@@ -26,4 +20,3 @@ class RandomBot(Agent):
         if not candidates:
             return Move.pass_turn()
         return Move.play(random.choice(candidates))
-# end::random_bot[]
diff --git a/agent/naive_fast.py b/agents/naive_fast.py
similarity index 86%
rename from agent/naive_fast.py
rename to agents/naive_fast.py
index 359dc10..5705c99 100644
--- a/agent/naive_fast.py
+++ b/agents/naive_fast.py
@@ -1,12 +1,11 @@
 import numpy as np
 
-from tugo.agent.base import Agent
-from tugo.agent.helpers_fast import is_point_an_eye
-from tugo.goboard import Move
-from tugo.gotypes import Point
+from tugo.agents.base import Agent
+from tugo.agents.helpers_fast import is_point_an_eye
+from game_logic.goboard import Move
+from game_logic.gotypes import Point
 
 
-__all__ = ['FastRandomBot']
 
 
 class FastRandomBot(Agent):
diff --git a/agents/policy_gradient_agent.py b/agents/policy_gradient_agent.py
new file mode 100644
index 0000000..3dcefd8
--- /dev/null
+++ b/agents/policy_gradient_agent.py
@@ -0,0 +1,91 @@
+"""Policy gradient learning."""
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch.optim import SGD
+
+from tugo.agents.base import Agent
+from tugo.agents.helpers import is_point_an_eye
+from game_logic import goboard
+
+
+def policy_gradient_loss(y_true, y_pred):
+    clip_pred = torch.clamp(y_pred, 1e-10, 1 - 1e-10)
+    loss = -1 * y_true * torch.log(clip_pred)
+    return torch.mean(torch.sum(loss, dim=1))
+
+
+class PolicyAgent(Agent):
+    """An agent that uses a deep policy network to select moves."""
+    def __init__(self, model, encoder):
+        super().__init__()
+        self._model = model
+        self._encoder = encoder
+        self._collector = None
+        self._temperature = 0.0
+
+    def predict(self, game_state):
+        encoded_state = self._encoder.encode(game_state)
+        input_tensor = torch.tensor([encoded_state], dtype=torch.float32).to('cuda')
+        with torch.no_grad():
+            output_tensor = self._model(input_tensor)
+        return output_tensor.cpu().numpy()[0]
+
+    def set_temperature(self, temperature):
+        self._temperature = temperature
+
+    def set_collector(self, collector):
+        self._collector = collector
+
+    def select_move(self, game_state):
+        num_moves = self._encoder.board_width * self._encoder.board_height
+
+        move_probs = self.predict(game_state)
+
+        move_probs = move_probs ** 3
+        eps = 1e-6
+        move_probs = np.clip(move_probs, eps, 1 - eps)
+        move_probs = move_probs / np.sum(move_probs)
+
+        candidates = np.arange(num_moves)
+        ranked_moves = np.random.choice(candidates, num_moves, replace=False, p=move_probs)
+
+        for point_idx in ranked_moves:
+            point = self._encoder.decode_point_index(point_idx)
+            if game_state.is_valid_move(goboard.Move.play(point)) and \
+                    not is_point_an_eye(game_state.board, point, game_state.next_player):
+                if self._collector is not None:
+                    self._collector.record_decision(state=board_tensor, action=point_idx)
+                return goboard.Move.play(point)
+
+        return goboard.Move.pass_turn()
+
+    def train(self, experience, lr=1e-7, clipnorm=1.0, batch_size=512):
+        opt = SGD(self._model.parameters(), lr=lr)
+        n = experience.states.shape[0]
+        num_moves = self._encoder.board_width * self._encoder.board_height
+        y = torch.zeros((n, num_moves))
+        for i in range(n):
+            action = experience.actions[i]
+            reward = experience.rewards[i]
+            y[i][action] = reward
+
+        for epoch in range(1):
+            permutation = torch.randperm(n)
+            for i in range(0, n, batch_size):
+                indices = permutation[i:i+batch_size]
+                batch_x, batch_y = experience.states[indices], y[indices]
+                opt.zero_grad()
+                outputs = self._model(batch_x)
+                loss = F.cross_entropy(outputs, batch_y)
+                loss.backward()
+                opt.step()
+
+    def save(self, path):
+        torch.save(self._model.state_dict(), path)
+
+    @classmethod
+    def load(cls, path, encoder):
+        model = AlphaGoModel(encoder.get_input_shape(), is_policy_net=True).to('cuda')
+        model.load_state_dict(torch.load(path))
+        return cls(model, encoder)
diff --git a/agent/predict.py b/agents/predict.py
similarity index 83%
rename from agent/predict.py
rename to agents/predict.py
index 51fd3aa..3b3bd27 100644
--- a/agent/predict.py
+++ b/agents/predict.py
@@ -1,16 +1,7 @@
-# tag::dl_agent_imports[]
-import numpy as np
-
-from tugo.agent.base import Agent
-from tugo.agent.helpers import is_point_an_eye
-from tugo import encoders
-from tugo import goboard
-from tugo import kerasutil
-from tugo.networks import AlphaGoModel
-# end::dl_agent_imports[]
-__all__ = [
-    'DeepLearningAgent',
-]
+from tugo.agents.base import Agent
+from tugo.agents.helpers import is_point_an_eye
+from game_logic import goboard
+from tugo.models import AlphaGoModel
 
 import numpy as np
 import torch
diff --git a/agent/termination.py b/agents/termination.py
similarity index 96%
rename from agent/termination.py
rename to agents/termination.py
index 14ef8d0..2083d64 100644
--- a/agent/termination.py
+++ b/agents/termination.py
@@ -1,7 +1,8 @@
 # tag::termination_imports[]
-from tugo import goboard
-from tugo.agent.base import Agent
-from tugo import scoring
+from game_logic import goboard, scoring
+from tugo.agents.base import Agent
+
+
 # end::termination_imports[]
 
 
diff --git a/alphago_policy_rl.py b/alphago_policy_rl.py
index 99e66a2..e1dfda1 100644
--- a/alphago_policy_rl.py
+++ b/alphago_policy_rl.py
@@ -1,8 +1,8 @@
 import torch
 from tqdm import tqdm
 
-from tugo.agent.pg import PolicyAgent
-from tugo.agent.predict import DeepLearningAgent
+from tugo.agents.policy_gradient_agent import PolicyAgent
+from tugo.agents.predict import DeepLearningAgent
 from tugo.encoders.alphago import AlphaGoEncoder
 from tugo.rl.simulate import experience_simulation
 
diff --git a/alphago_policy_sl.py b/alphago_policy_sl.py
index 6dafcf8..668cb27 100644
--- a/alphago_policy_sl.py
+++ b/alphago_policy_sl.py
@@ -1,7 +1,7 @@
-from tugo.data.parallel_processor import GoDataProcessor
+from tugo.data_processing.parallel_processor import GoDataProcessor
 from tugo.encoders.alphago import AlphaGoEncoder
-from tugo.agent.predict import DeepLearningAgent
-from tugo.networks.alphago import AlphaGoModel
+from tugo.agents.predict import DeepLearningAgent
+from tugo.models.alphago import AlphaGoModel
 
 import torch
 from torch import nn
diff --git a/data/__init__.py b/data_processing/__init__.py
similarity index 100%
rename from data/__init__.py
rename to data_processing/__init__.py
diff --git a/data/generator.py b/data_processing/generator.py
similarity index 100%
rename from data/generator.py
rename to data_processing/generator.py
diff --git a/data/index_processor.py b/data_processing/index_processor.py
similarity index 100%
rename from data/index_processor.py
rename to data_processing/index_processor.py
diff --git a/data/parallel_processor.py b/data_processing/parallel_processor.py
similarity index 97%
rename from data/parallel_processor.py
rename to data_processing/parallel_processor.py
index 101e0a1..ddfa1ff 100644
--- a/data/parallel_processor.py
+++ b/data_processing/parallel_processor.py
@@ -12,11 +12,11 @@ from os import sys
 import torch
 
 from tugo.gosgf import Sgf_game
-from tugo.goboard_fast import Board, GameState, Move
-from tugo.gotypes import Player, Point
-from tugo.data.index_processor import KGSIndex
-from tugo.data.sampling import Sampler
-from tugo.data.generator import DataGenerator
+from game_logic.goboard_fast import Board, GameState, Move
+from game_logic.gotypes import Player, Point
+from tugo.data_processing.index_processor import KGSIndex
+from tugo.data_processing.sampling import Sampler
+from tugo.data_processing.generator import DataGenerator
 from tugo.encoders.base import get_encoder_by_name
 from torch.utils.data import TensorDataset
 
diff --git a/data/sampling.py b/data_processing/sampling.py
similarity index 98%
rename from data/sampling.py
rename to data_processing/sampling.py
index 1231af6..853d59f 100644
--- a/data/sampling.py
+++ b/data_processing/sampling.py
@@ -5,7 +5,7 @@ from __future__ import print_function
 from __future__ import absolute_import
 import os
 import random
-from tugo.data.index_processor import KGSIndex
+from tugo.data_processing.index_processor import KGSIndex
 from six.moves import range
 
 
diff --git a/encoders/alphago.py b/encoders/alphago.py
index aa964af..4c9f7dc 100644
--- a/encoders/alphago.py
+++ b/encoders/alphago.py
@@ -1,9 +1,8 @@
 from tugo.encoders.base import Encoder
 from tugo.encoders.utils import is_ladder_escape, is_ladder_capture
-from tugo.gotypes import Point, Player
-from tugo.goboard_fast import Move
-from tugo.agent.helpers_fast import is_point_an_eye
-import numpy as np
+from game_logic.gotypes import Point, Player
+from game_logic.goboard_fast import Move
+from tugo.agents.helpers_fast import is_point_an_eye
 import torch
 
 """
diff --git a/encoders/oneplane.py b/encoders/oneplane.py
index 23ac364..37449fa 100644
--- a/encoders/oneplane.py
+++ b/encoders/oneplane.py
@@ -1,7 +1,7 @@
 import numpy as np
 
 from tugo.encoders.base import Encoder
-from tugo.goboard import Point
+from game_logic.goboard import Point
 
 
 
diff --git a/encoders/sevenplane.py b/encoders/sevenplane.py
index e817312..f036f58 100644
--- a/encoders/sevenplane.py
+++ b/encoders/sevenplane.py
@@ -2,7 +2,7 @@
 import numpy as np
 
 from tugo.encoders.base import Encoder
-from tugo.goboard import Move, Point
+from game_logic.goboard import Move, Point
 
 
 class SevenPlaneEncoder(Encoder):
diff --git a/encoders/simple.py b/encoders/simple.py
index ebb2906..307f4ea 100644
--- a/encoders/simple.py
+++ b/encoders/simple.py
@@ -1,8 +1,8 @@
 import numpy as np
 
 from tugo.encoders.base import Encoder
-from tugo.goboard import Move
-from tugo.gotypes import Player, Point
+from game_logic.goboard import Move
+from game_logic.gotypes import Player, Point
 
 
 class SimpleEncoder(Encoder):
diff --git a/encoders/utils.py b/encoders/utils.py
index 609106a..bf29b65 100644
--- a/encoders/utils.py
+++ b/encoders/utils.py
@@ -1,4 +1,4 @@
-from tugo.goboard import Move
+from game_logic.goboard import Move
 
 
 def is_ladder_capture(game_state, candidate, recursion_depth=50):
diff --git a/game_logic/__init__.py b/game_logic/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/goboard.py b/game_logic/goboard.py
similarity index 96%
rename from goboard.py
rename to game_logic/goboard.py
index 193e189..dc84d92 100644
--- a/goboard.py
+++ b/game_logic/goboard.py
@@ -1,15 +1,7 @@
 import copy
-from tugo.gotypes import Player, Point
-from tugo.scoring import compute_game_result
-# tag::import_zobrist[]
-from tugo import zobrist
-# end::import_zobrist[]
-
-__all__ = [
-    'Board',
-    'GameState',
-    'Move',
-]
+from game_logic.gotypes import Player, Point
+from game_logic.scoring import compute_game_result
+from game_logic import zobrist_hash
 
 
 class IllegalMoveError(Exception):
@@ -64,7 +56,7 @@ class Board:
         self.num_rows = num_rows
         self.num_cols = num_cols
         self._grid = {}
-        self._hash = zobrist.EMPTY_BOARD
+        self._hash = zobrist_hash.EMPTY_BOARD
 # end::init_zobrist[]
 
     def place_stone(self, player, point):
@@ -97,7 +89,7 @@ class Board:
         for new_string_point in new_string.stones:
             self._grid[new_string_point] = new_string
 
-        self._hash ^= zobrist.HASH_CODE[point, player]  # <3>
+        self._hash ^= zobrist_hash.HASH_CODE[point, player]  # <3>
 
         for other_color_string in adjacent_opposite_color:
             replacement = other_color_string.without_liberty(point)  # <4>
@@ -128,7 +120,7 @@ class Board:
                     self._replace_string(neighbor_string.with_liberty(point))
             self._grid[point] = None
 
-            self._hash ^= zobrist.HASH_CODE[point, string.color]  # <3>
+            self._hash ^= zobrist_hash.HASH_CODE[point, string.color]  # <3>
 # <1> This new helper method updates our Go board grid.
 # <2> Removing a string can create liberties for other strings.
 # <3> With Zobrist hashing, you need to unapply the hash for this move.
diff --git a/goboard_fast.py b/game_logic/goboard_fast.py
similarity index 96%
rename from goboard_fast.py
rename to game_logic/goboard_fast.py
index e36603d..11d68c0 100644
--- a/goboard_fast.py
+++ b/game_logic/goboard_fast.py
@@ -1,14 +1,8 @@
 import copy
-from tugo.gotypes import Player, Point
-from tugo.scoring import compute_game_result
-from tugo import zobrist
-from tugo.utils import MoveAge
-
-__all__ = [
-    'Board',
-    'GameState',
-    'Move',
-]
+from game_logic.gotypes import Player, Point
+from game_logic.scoring import compute_game_result
+from game_logic import zobrist_hash
+from tugo.print_utils import MoveAge
 
 neighbor_tables = {}
 corner_tables = {}
@@ -96,7 +90,7 @@ class Board():
         self.num_rows = num_rows
         self.num_cols = num_cols
         self._grid = {}
-        self._hash = zobrist.EMPTY_BOARD
+        self._hash = zobrist_hash.EMPTY_BOARD
 
         global neighbor_tables
         dim = (num_rows, num_cols)
@@ -144,9 +138,9 @@ class Board():
         for new_string_point in new_string.stones:
             self._grid[new_string_point] = new_string
         # Remove empty-point hash code.
-        self._hash ^= zobrist.HASH_CODE[point, None]
+        self._hash ^= zobrist_hash.HASH_CODE[point, None]
         # Add filled point hash code.
-        self._hash ^= zobrist.HASH_CODE[point, player]
+        self._hash ^= zobrist_hash.HASH_CODE[point, player]
 # end::apply_zobrist[]
 
         # 2. Reduce liberties of any adjacent strings of the opposite
@@ -176,9 +170,9 @@ class Board():
                     self._replace_string(neighbor_string.with_liberty(point))
             self._grid[point] = None
             # Remove filled point hash code.
-            self._hash ^= zobrist.HASH_CODE[point, string.color]
+            self._hash ^= zobrist_hash.HASH_CODE[point, string.color]
             # Add empty point hash code.
-            self._hash ^= zobrist.HASH_CODE[point, None]
+            self._hash ^= zobrist_hash.HASH_CODE[point, None]
 
     def is_self_capture(self, player, point):
         friendly_strings = []
diff --git a/gotypes.py b/game_logic/gotypes.py
similarity index 84%
rename from gotypes.py
rename to game_logic/gotypes.py
index f3af4cf..c072395 100644
--- a/gotypes.py
+++ b/game_logic/gotypes.py
@@ -1,13 +1,6 @@
-# tag::enumimport[]
 import enum
-# end::enumimport[]
-# tag::namedtuple[]
 from collections import namedtuple
-# end::namedtuple[]
-__all__ = [
-    'Player',
-    'Point',
-]
+
 
 
 # tag::color[]
diff --git a/scoring.py b/game_logic/scoring.py
similarity index 98%
rename from scoring.py
rename to game_logic/scoring.py
index 3386d4c..df52e9f 100644
--- a/scoring.py
+++ b/game_logic/scoring.py
@@ -2,8 +2,8 @@
 from __future__ import absolute_import
 from collections import namedtuple
 
-from tugo.gotypes import Player, Point
-# end::scoring_imports[]
+from game_logic.gotypes import Player, Point
+
 
 
 # tag::scoring_territory[]
diff --git a/zobrist.py b/game_logic/zobrist_hash.py
similarity index 99%
rename from zobrist.py
rename to game_logic/zobrist_hash.py
index c699d8b..ccdef9f 100644
--- a/zobrist.py
+++ b/game_logic/zobrist_hash.py
@@ -1,6 +1,5 @@
-from tugo.gotypes import Player, Point
+from game_logic.gotypes import Player, Point
 
-__all__ = ['HASH_CODE', 'EMPTY_BOARD']
 
 HASH_CODE = {
     (Point(row=1, col=1), None): 6402364705153495313,
diff --git a/go_sgf/sgf.py b/go_sgf/sgf.py
index a25b30a..ffa7e84 100644
--- a/go_sgf/sgf.py
+++ b/go_sgf/sgf.py
@@ -13,11 +13,6 @@ import six
 from . import sgf_grammar
 from . import sgf_properties
 
-__all__ = [
-    'Node',
-    'Sgf_game',
-    'Tree_node',
-]
 
 
 class Node:
diff --git a/kerasutil.py b/kerasutil.py
deleted file mode 100644
index 83c9125..0000000
--- a/kerasutil.py
+++ /dev/null
@@ -1,30 +0,0 @@
-import tempfile
-import os
-import torch
-
-def save_model_to_hdf5_group(model, f):
-    tempfd, tempfname = tempfile.mkstemp(prefix='tmp-torchmodel')
-    try:
-        os.close(tempfd)
-        torch.save(model, tempfname)
-        with open(tempfname, 'rb') as model_file:
-            model_data = model_file.read()
-        f.create_dataset('torchmodel', data=model_data)
-    finally:
-        os.unlink(tempfname)
-
-def load_model_from_hdf5_group(f, map_location=None):
-    tempfd, tempfname = tempfile.mkstemp(prefix='tmp-torchmodel')
-    try:
-        os.close(tempfd)
-        with open(tempfname, 'wb') as model_file:
-            model_file.write(f['torchmodel'][()])
-        model = torch.load(tempfname, map_location=map_location)
-        return model
-    finally:
-        os.unlink(tempfname)
-
-def set_gpu_memory_target(device, frac):
-    # This function is not required in PyTorch since it doesn't pre-allocate all GPU memory by default.
-    pass
-
diff --git a/networks/__init__.py b/models/__init__.py
similarity index 100%
rename from networks/__init__.py
rename to models/__init__.py
diff --git a/networks/alphago.py b/models/alphago.py
similarity index 100%
rename from networks/alphago.py
rename to models/alphago.py
diff --git a/networks/alphago_zero.py b/models/alphago_zero.py
similarity index 100%
rename from networks/alphago_zero.py
rename to models/alphago_zero.py
diff --git a/utils.py b/print_utils.py
similarity index 97%
rename from utils.py
rename to print_utils.py
index 815d0b0..c4fa1c9 100644
--- a/utils.py
+++ b/print_utils.py
@@ -3,8 +3,7 @@ import subprocess
 
 import numpy as np
 
-# tag::print_utils[]
-from tugo import gotypes
+from game_logic import gotypes
 
 COLS = 'ABCDEFGHJKLMNOPQRST'
 STONE_TO_CHAR = {
diff --git a/tree_search/__init__.py b/tree_search/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/web/__init__.py b/web/__init__.py
new file mode 100644
index 0000000..e69de29