diff --git a/agents/policy_gradient_agent.py b/agents/policy_gradient_agent.py
index 8e80991..e580fc8 100644
--- a/agents/policy_gradient_agent.py
+++ b/agents/policy_gradient_agent.py
@@ -6,7 +6,8 @@ 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
+# from tugo.game_logic import goboard
+from tugo.game_logic import goboard_fast as goboard
 
 
 def policy_gradient_loss(y_true, y_pred):
@@ -26,7 +27,8 @@ class PolicyAgent(Agent):
 
     def predict(self, game_state):
         encoded_state = self._encoder.encode(game_state)
-        input_tensor = torch.tensor([encoded_state], dtype=torch.float32).to('cuda')
+        # input_tensor = torch.tensor([encoded_state], dtype=torch.float32).to('cuda')
+        input_tensor = encoded_state.float().to('cuda')
         with torch.no_grad():
             output_tensor = self._model(input_tensor)
         return output_tensor.cpu().numpy()[0]
diff --git a/agents/predict.py b/agents/predict.py
index 83ea33f..30d3d06 100644
--- a/agents/predict.py
+++ b/agents/predict.py
@@ -1,6 +1,7 @@
 from tugo.agents.base import Agent
 from tugo.agents.helpers import is_point_an_eye
-from game_logic import goboard
+# from game_logic import goboard
+from tugo.game_logic import goboard_fast as goboard
 from tugo.models import AlphaGoModel
 
 import numpy as np
@@ -44,6 +45,7 @@ class DeepLearningAgent(Agent):
     @classmethod
     def load(cls, file_path, encoder):
         device = 'cuda' if torch.cuda.is_available() else 'cpu'
-        model = AlphaGoModel(encoder.get_input_shape(), is_policy_net=True).to(device)
+        # model = AlphaGoModel(encoder.get_input_shape(), is_policy_net=True).to(device)
+        model = AlphaGoModel(encoder.shape(), is_policy_net=True, num_classes=361).to(device)
         model.load_state_dict(torch.load(file_path))
         return cls(model, encoder)
diff --git a/encoders/utils.py b/encoders/utils.py
index bf29b65..e64bb97 100644
--- a/encoders/utils.py
+++ b/encoders/utils.py
@@ -1,4 +1,4 @@
-from game_logic.goboard import Move
+from tugo.game_logic.goboard_fast import Move
 
 
 def is_ladder_capture(game_state, candidate, recursion_depth=50):
diff --git a/game_logic/goboard.py b/game_logic/goboard.py
index dc84d92..05cbf38 100644
--- a/game_logic/goboard.py
+++ b/game_logic/goboard.py
@@ -1,299 +1,299 @@
-import copy
-from game_logic.gotypes import Player, Point
-from game_logic.scoring import compute_game_result
-from game_logic import zobrist_hash
-
-
-class IllegalMoveError(Exception):
-    pass
-
-
-# tag::fast_go_strings[]
-class GoString:
-    def __init__(self, color, stones, liberties):
-        self.color = color
-        self.stones = frozenset(stones)
-        self.liberties = frozenset(liberties)  # <1>
-
-    def without_liberty(self, point):  # <2>
-        new_liberties = self.liberties - set([point])
-        return GoString(self.color, self.stones, new_liberties)
-
-    def with_liberty(self, point):
-        new_liberties = self.liberties | set([point])
-        return GoString(self.color, self.stones, new_liberties)
-# <1> `stones` and `liberties` are now immutable `frozenset` instances
-# <2> The `without_liberty` methods replaces the previous `remove_liberty` method...
-# <3> ... and `with_liberty` replaces `add_liberty`.
-# end::fast_go_strings[]
-
-    def merged_with(self, string):
-        """Return a new string containing all stones in both strings."""
-        assert string.color == self.color
-        combined_stones = self.stones | string.stones
-        return GoString(
-            self.color,
-            combined_stones,
-            (self.liberties | string.liberties) - combined_stones)
-
-    @property
-    def num_liberties(self):
-        return len(self.liberties)
-
-    def __eq__(self, other):
-        return isinstance(other, GoString) and \
-            self.color == other.color and \
-            self.stones == other.stones and \
-            self.liberties == other.liberties
-
-    def __deepcopy__(self, memodict={}):
-        return GoString(self.color, self.stones, copy.deepcopy(self.liberties))
-
-
-# tag::init_zobrist[]
-class Board:
-    def __init__(self, num_rows, num_cols):
-        self.num_rows = num_rows
-        self.num_cols = num_cols
-        self._grid = {}
-        self._hash = zobrist_hash.EMPTY_BOARD
-# end::init_zobrist[]
-
-    def place_stone(self, player, point):
-        assert self.is_on_grid(point)
-        if self._grid.get(point) is not None:
-            print('Illegal play on %s' % str(point))
-        assert self._grid.get(point) is None
-        # 0. Examine the adjacent points.
-        adjacent_same_color = []
-        adjacent_opposite_color = []
-        liberties = []
-        for neighbor in point.neighbors():
-            if not self.is_on_grid(neighbor):
-                continue
-            neighbor_string = self._grid.get(neighbor)
-            if neighbor_string is None:
-                liberties.append(neighbor)
-            elif neighbor_string.color == player:
-                if neighbor_string not in adjacent_same_color:
-                    adjacent_same_color.append(neighbor_string)
-            else:
-                if neighbor_string not in adjacent_opposite_color:
-                    adjacent_opposite_color.append(neighbor_string)
-        new_string = GoString(player, [point], liberties)
-# tag::apply_zobrist[]
-        new_string = GoString(player, [point], liberties)  # <1>
-
-        for same_color_string in adjacent_same_color:  # <2>
-            new_string = new_string.merged_with(same_color_string)
-        for new_string_point in new_string.stones:
-            self._grid[new_string_point] = new_string
-
-        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>
-            if replacement.num_liberties:
-                self._replace_string(other_color_string.without_liberty(point))
-            else:
-                self._remove_string(other_color_string)  # <5>
-# <1> Until this line `place_stone` remains the same.
-# <2> You merge any adjacent strings of the same color.
-# <3> Next, you apply the hash code for this point and player
-# <4> Then you reduce liberties of any adjacent strings of the opposite color.
-# <5> If any opposite color strings now have zero liberties, remove them.
-# end::apply_zobrist[]
-
-
-# tag::unapply_zobrist[]
-    def _replace_string(self, new_string):  # <1>
-        for point in new_string.stones:
-            self._grid[point] = new_string
-
-    def _remove_string(self, string):
-        for point in string.stones:
-            for neighbor in point.neighbors():  # <2>
-                neighbor_string = self._grid.get(neighbor)
-                if neighbor_string is None:
-                    continue
-                if neighbor_string is not string:
-                    self._replace_string(neighbor_string.with_liberty(point))
-            self._grid[point] = None
-
-            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.
-# end::unapply_zobrist[]
-
-    def is_on_grid(self, point):
-        return 1 <= point.row <= self.num_rows and \
-            1 <= point.col <= self.num_cols
-
-    def get(self, point):
-        """Return the content of a point on the board.
-
-        Returns None if the point is empty, or a Player if there is a
-        stone on that point.
-        """
-        string = self._grid.get(point)
-        if string is None:
-            return None
-        return string.color
-
-    def get_go_string(self, point):
-        """Return the entire string of stones at a point.
-
-        Returns None if the point is empty, or a GoString if there is
-        a stone on that point.
-        """
-        string = self._grid.get(point)
-        if string is None:
-            return None
-        return string
-
-    def __eq__(self, other):
-        return isinstance(other, Board) and \
-            self.num_rows == other.num_rows and \
-            self.num_cols == other.num_cols and \
-            self._hash() == other._hash()
-
-    def __deepcopy__(self, memodict={}):
-        copied = Board(self.num_rows, self.num_cols)
-        # Can do a shallow copy b/c the dictionary maps tuples
-        # (immutable) to GoStrings (also immutable)
-        copied._grid = copy.copy(self._grid)
-        copied._hash = self._hash
-        return copied
-
-# tag::return_zobrist[]
-    def zobrist_hash(self):
-        return self._hash
-# end::return_zobrist[]
-
-
-class Move:
-    """Any action a player can play on a turn.
-    Exactly one of is_play, is_pass, is_resign will be set.
-    """
-    def __init__(self, point=None, is_pass=False, is_resign=False):
-        assert (point is not None) ^ is_pass ^ is_resign
-        self.point = point
-        self.is_play = (self.point is not None)
-        self.is_pass = is_pass
-        self.is_resign = is_resign
-
-    @classmethod
-    def play(cls, point):
-        """A move that places a stone on the board."""
-        return Move(point=point)
-
-    @classmethod
-    def pass_turn(cls):
-        return Move(is_pass=True)
-
-    @classmethod
-    def resign(cls):
-        return Move(is_resign=True)
-
-    def __str__(self):
-        if self.is_pass:
-            return 'pass'
-        if self.is_resign:
-            return 'resign'
-        return '(r %d, c %d)' % (self.point.row, self.point.col)
-
-
-# tag::init_state_zobrist[]
-class GameState:
-    def __init__(self, board, next_player, previous, move):
-        self.board = board
-        self.next_player = next_player
-        self.previous_state = previous
-        if self.previous_state is None:
-            self.previous_states = frozenset()
-        else:
-            self.previous_states = frozenset(
-                previous.previous_states |
-                {(previous.next_player, previous.board.zobrist_hash())})
-        self.last_move = move
-# end::init_state_zobrist[]
-
-    def apply_move(self, move):
-        """Return the new GameState after applying the move."""
-        if move.is_play:
-            next_board = copy.deepcopy(self.board)
-            next_board.place_stone(self.next_player, move.point)
-        else:
-            next_board = self.board
-        return GameState(next_board, self.next_player.other, self, move)
-
-    @classmethod
-    def new_game(cls, board_size):
-        if isinstance(board_size, int):
-            board_size = (board_size, board_size)
-        board = Board(*board_size)
-        return GameState(board, Player.black, None, None)
-
-    def is_move_self_capture(self, player, move):
-        if not move.is_play:
-            return False
-        next_board = copy.deepcopy(self.board)
-        next_board.place_stone(player, move.point)
-        new_string = next_board.get_go_string(move.point)
-        return new_string.num_liberties == 0
-
-    @property
-    def situation(self):
-        return (self.next_player, self.board)
-
-# tag::ko_zobrist[]
-    def does_move_violate_ko(self, player, move):
-        if not move.is_play:
-            return False
-        next_board = copy.deepcopy(self.board)
-        next_board.place_stone(player, move.point)
-        next_situation = (player.other, next_board.zobrist_hash())
-        return next_situation in self.previous_states
-# end::ko_zobrist[]
-
-    def is_valid_move(self, move):
-        if self.is_over():
-            return False
-        if move.is_pass or move.is_resign:
-            return True
-        return (
-            self.board.get(move.point) is None and
-            not self.is_move_self_capture(self.next_player, move) and
-            not self.does_move_violate_ko(self.next_player, move))
-
-    def is_over(self):
-        if self.last_move is None:
-            return False
-        if self.last_move.is_resign:
-            return True
-        second_last_move = self.previous_state.last_move
-        if second_last_move is None:
-            return False
-        return self.last_move.is_pass and second_last_move.is_pass
-
-    def legal_moves(self):
-        moves = []
-        for row in range(1, self.board.num_rows + 1):
-            for col in range(1, self.board.num_cols + 1):
-                move = Move.play(Point(row, col))
-                if self.is_valid_move(move):
-                    moves.append(move)
-        # These two moves are always legal.
-        moves.append(Move.pass_turn())
-        moves.append(Move.resign())
-
-        return moves
-
-    def winner(self):
-        if not self.is_over():
-            return None
-        if self.last_move.is_resign:
-            return self.next_player
-        game_result = compute_game_result(self)
-        return game_result.winner
+# import copy
+# from tugo.game_logic.gotypes import Player, Point
+# from tugo.game_logic.scoring import compute_game_result
+# from tugo.game_logic import zobrist_hash
+#
+#
+# class IllegalMoveError(Exception):
+#     pass
+#
+#
+# # tag::fast_go_strings[]
+# class GoString:
+#     def __init__(self, color, stones, liberties):
+#         self.color = color
+#         self.stones = frozenset(stones)
+#         self.liberties = frozenset(liberties)  # <1>
+#
+#     def without_liberty(self, point):  # <2>
+#         new_liberties = self.liberties - set([point])
+#         return GoString(self.color, self.stones, new_liberties)
+#
+#     def with_liberty(self, point):
+#         new_liberties = self.liberties | set([point])
+#         return GoString(self.color, self.stones, new_liberties)
+# # <1> `stones` and `liberties` are now immutable `frozenset` instances
+# # <2> The `without_liberty` methods replaces the previous `remove_liberty` method...
+# # <3> ... and `with_liberty` replaces `add_liberty`.
+# # end::fast_go_strings[]
+#
+#     def merged_with(self, string):
+#         """Return a new string containing all stones in both strings."""
+#         assert string.color == self.color
+#         combined_stones = self.stones | string.stones
+#         return GoString(
+#             self.color,
+#             combined_stones,
+#             (self.liberties | string.liberties) - combined_stones)
+#
+#     @property
+#     def num_liberties(self):
+#         return len(self.liberties)
+#
+#     def __eq__(self, other):
+#         return isinstance(other, GoString) and \
+#             self.color == other.color and \
+#             self.stones == other.stones and \
+#             self.liberties == other.liberties
+#
+#     def __deepcopy__(self, memodict={}):
+#         return GoString(self.color, self.stones, copy.deepcopy(self.liberties))
+#
+#
+# # tag::init_zobrist[]
+# class Board:
+#     def __init__(self, num_rows, num_cols):
+#         self.num_rows = num_rows
+#         self.num_cols = num_cols
+#         self._grid = {}
+#         self._hash = zobrist_hash.EMPTY_BOARD
+# # end::init_zobrist[]
+#
+#     def place_stone(self, player, point):
+#         assert self.is_on_grid(point)
+#         if self._grid.get(point) is not None:
+#             print('Illegal play on %s' % str(point))
+#         assert self._grid.get(point) is None
+#         # 0. Examine the adjacent points.
+#         adjacent_same_color = []
+#         adjacent_opposite_color = []
+#         liberties = []
+#         for neighbor in point.neighbors():
+#             if not self.is_on_grid(neighbor):
+#                 continue
+#             neighbor_string = self._grid.get(neighbor)
+#             if neighbor_string is None:
+#                 liberties.append(neighbor)
+#             elif neighbor_string.color == player:
+#                 if neighbor_string not in adjacent_same_color:
+#                     adjacent_same_color.append(neighbor_string)
+#             else:
+#                 if neighbor_string not in adjacent_opposite_color:
+#                     adjacent_opposite_color.append(neighbor_string)
+#         new_string = GoString(player, [point], liberties)
+# # tag::apply_zobrist[]
+#         new_string = GoString(player, [point], liberties)  # <1>
+#
+#         for same_color_string in adjacent_same_color:  # <2>
+#             new_string = new_string.merged_with(same_color_string)
+#         for new_string_point in new_string.stones:
+#             self._grid[new_string_point] = new_string
+#
+#         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>
+#             if replacement.num_liberties:
+#                 self._replace_string(other_color_string.without_liberty(point))
+#             else:
+#                 self._remove_string(other_color_string)  # <5>
+# # <1> Until this line `place_stone` remains the same.
+# # <2> You merge any adjacent strings of the same color.
+# # <3> Next, you apply the hash code for this point and player
+# # <4> Then you reduce liberties of any adjacent strings of the opposite color.
+# # <5> If any opposite color strings now have zero liberties, remove them.
+# # end::apply_zobrist[]
+#
+#
+# # tag::unapply_zobrist[]
+#     def _replace_string(self, new_string):  # <1>
+#         for point in new_string.stones:
+#             self._grid[point] = new_string
+#
+#     def _remove_string(self, string):
+#         for point in string.stones:
+#             for neighbor in point.neighbors():  # <2>
+#                 neighbor_string = self._grid.get(neighbor)
+#                 if neighbor_string is None:
+#                     continue
+#                 if neighbor_string is not string:
+#                     self._replace_string(neighbor_string.with_liberty(point))
+#             self._grid[point] = None
+#
+#             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.
+# # end::unapply_zobrist[]
+#
+#     def is_on_grid(self, point):
+#         return 1 <= point.row <= self.num_rows and \
+#             1 <= point.col <= self.num_cols
+#
+#     def get(self, point):
+#         """Return the content of a point on the board.
+#
+#         Returns None if the point is empty, or a Player if there is a
+#         stone on that point.
+#         """
+#         string = self._grid.get(point)
+#         if string is None:
+#             return None
+#         return string.color
+#
+#     def get_go_string(self, point):
+#         """Return the entire string of stones at a point.
+#
+#         Returns None if the point is empty, or a GoString if there is
+#         a stone on that point.
+#         """
+#         string = self._grid.get(point)
+#         if string is None:
+#             return None
+#         return string
+#
+#     def __eq__(self, other):
+#         return isinstance(other, Board) and \
+#             self.num_rows == other.num_rows and \
+#             self.num_cols == other.num_cols and \
+#             self._hash() == other._hash()
+#
+#     def __deepcopy__(self, memodict={}):
+#         copied = Board(self.num_rows, self.num_cols)
+#         # Can do a shallow copy b/c the dictionary maps tuples
+#         # (immutable) to GoStrings (also immutable)
+#         copied._grid = copy.copy(self._grid)
+#         copied._hash = self._hash
+#         return copied
+#
+# # tag::return_zobrist[]
+#     def zobrist_hash(self):
+#         return self._hash
+# # end::return_zobrist[]
+#
+#
+# class Move:
+#     """Any action a player can play on a turn.
+#     Exactly one of is_play, is_pass, is_resign will be set.
+#     """
+#     def __init__(self, point=None, is_pass=False, is_resign=False):
+#         assert (point is not None) ^ is_pass ^ is_resign
+#         self.point = point
+#         self.is_play = (self.point is not None)
+#         self.is_pass = is_pass
+#         self.is_resign = is_resign
+#
+#     @classmethod
+#     def play(cls, point):
+#         """A move that places a stone on the board."""
+#         return Move(point=point)
+#
+#     @classmethod
+#     def pass_turn(cls):
+#         return Move(is_pass=True)
+#
+#     @classmethod
+#     def resign(cls):
+#         return Move(is_resign=True)
+#
+#     def __str__(self):
+#         if self.is_pass:
+#             return 'pass'
+#         if self.is_resign:
+#             return 'resign'
+#         return '(r %d, c %d)' % (self.point.row, self.point.col)
+#
+#
+# # tag::init_state_zobrist[]
+# class GameState:
+#     def __init__(self, board, next_player, previous, move):
+#         self.board = board
+#         self.next_player = next_player
+#         self.previous_state = previous
+#         if self.previous_state is None:
+#             self.previous_states = frozenset()
+#         else:
+#             self.previous_states = frozenset(
+#                 previous.previous_states |
+#                 {(previous.next_player, previous.board.zobrist_hash())})
+#         self.last_move = move
+# # end::init_state_zobrist[]
+#
+#     def apply_move(self, move):
+#         """Return the new GameState after applying the move."""
+#         if move.is_play:
+#             next_board = copy.deepcopy(self.board)
+#             next_board.place_stone(self.next_player, move.point)
+#         else:
+#             next_board = self.board
+#         return GameState(next_board, self.next_player.other, self, move)
+#
+#     @classmethod
+#     def new_game(cls, board_size):
+#         if isinstance(board_size, int):
+#             board_size = (board_size, board_size)
+#         board = Board(*board_size)
+#         return GameState(board, Player.black, None, None)
+#
+#     def is_move_self_capture(self, player, move):
+#         if not move.is_play:
+#             return False
+#         next_board = copy.deepcopy(self.board)
+#         next_board.place_stone(player, move.point)
+#         new_string = next_board.get_go_string(move.point)
+#         return new_string.num_liberties == 0
+#
+#     @property
+#     def situation(self):
+#         return (self.next_player, self.board)
+#
+# # tag::ko_zobrist[]
+#     def does_move_violate_ko(self, player, move):
+#         if not move.is_play:
+#             return False
+#         next_board = copy.deepcopy(self.board)
+#         next_board.place_stone(player, move.point)
+#         next_situation = (player.other, next_board.zobrist_hash())
+#         return next_situation in self.previous_states
+# # end::ko_zobrist[]
+#
+#     def is_valid_move(self, move):
+#         if self.is_over():
+#             return False
+#         if move.is_pass or move.is_resign:
+#             return True
+#         return (
+#             self.board.get(move.point) is None and
+#             not self.is_move_self_capture(self.next_player, move) and
+#             not self.does_move_violate_ko(self.next_player, move))
+#
+#     def is_over(self):
+#         if self.last_move is None:
+#             return False
+#         if self.last_move.is_resign:
+#             return True
+#         second_last_move = self.previous_state.last_move
+#         if second_last_move is None:
+#             return False
+#         return self.last_move.is_pass and second_last_move.is_pass
+#
+#     def legal_moves(self):
+#         moves = []
+#         for row in range(1, self.board.num_rows + 1):
+#             for col in range(1, self.board.num_cols + 1):
+#                 move = Move.play(Point(row, col))
+#                 if self.is_valid_move(move):
+#                     moves.append(move)
+#         # These two moves are always legal.
+#         moves.append(Move.pass_turn())
+#         moves.append(Move.resign())
+#
+#         return moves
+#
+#     def winner(self):
+#         if not self.is_over():
+#             return None
+#         if self.last_move.is_resign:
+#             return self.next_player
+#         game_result = compute_game_result(self)
+#         return game_result.winner
diff --git a/game_logic/goboard_fast.py b/game_logic/goboard_fast.py
index 11d68c0..19f2dca 100644
--- a/game_logic/goboard_fast.py
+++ b/game_logic/goboard_fast.py
@@ -1,7 +1,7 @@
 import copy
-from game_logic.gotypes import Player, Point
-from game_logic.scoring import compute_game_result
-from game_logic import zobrist_hash
+from tugo.game_logic.gotypes import Player, Point
+from tugo.game_logic.scoring import compute_game_result
+from tugo.game_logic import zobrist_hash
 from tugo.print_utils import MoveAge
 
 neighbor_tables = {}
diff --git a/game_logic/scoring.py b/game_logic/scoring.py
index df52e9f..6cd89b2 100644
--- a/game_logic/scoring.py
+++ b/game_logic/scoring.py
@@ -2,7 +2,7 @@
 from __future__ import absolute_import
 from collections import namedtuple
 
-from game_logic.gotypes import Player, Point
+from tugo.game_logic.gotypes import Player, Point
 
 
 
diff --git a/game_logic/zobrist_hash.py b/game_logic/zobrist_hash.py
index ccdef9f..10de82f 100644
--- a/game_logic/zobrist_hash.py
+++ b/game_logic/zobrist_hash.py
@@ -1,4 +1,4 @@
-from game_logic.gotypes import Player, Point
+from tugo.game_logic.gotypes import Player, Point
 
 
 HASH_CODE = {
diff --git a/rl/ac.py b/rl/ac.py
index ca3ddcf..6f99b63 100644
--- a/rl/ac.py
+++ b/rl/ac.py
@@ -3,7 +3,8 @@ import torch
 import torch.optim as optim
 
 from tugo import encoders
-from tugo.game_logic import goboard
+# from tugo.game_logic import goboard
+from tugo.game_logic import goboard_fast as goboard
 from tugo.agents import Agent
 from tugo.agents.helpers import is_point_an_eye
 
diff --git a/rl/ac_pass.py b/rl/ac_pass.py
index 50fe50d..ed65061 100644
--- a/rl/ac_pass.py
+++ b/rl/ac_pass.py
@@ -2,7 +2,8 @@ import torch
 import numpy as np
 
 from tugo import encoders
-from tugo.game_logic import goboard
+# from tugo.game_logic import goboard
+from tugo.game_logic import goboard_fast as goboard
 from tugo.agents import Agent
 
 
diff --git a/rl/q.py b/rl/q.py
index d2629ab..9f397d8 100644
--- a/rl/q.py
+++ b/rl/q.py
@@ -4,7 +4,8 @@ import torch.optim as optim
 from torch.distributions import Categorical
 
 from tugo import encoders
-from tugo.game_logic import goboard
+# from tugo.game_logic import goboard
+from tugo.game_logic import goboard_fast as goboard
 from tugo.agents import Agent
 from tugo.agents.helpers import is_point_an_eye
 
diff --git a/rl/simulate.py b/rl/simulate.py
index b1cacc2..d3adf65 100644
--- a/rl/simulate.py
+++ b/rl/simulate.py
@@ -1,4 +1,4 @@
-from tugo import rl
+from tugo.rl import experience
 from tugo.game_logic import scoring
 from tugo.game_logic import goboard_fast as goboard
 from tugo.game_logic.gotypes import Player
@@ -32,8 +32,8 @@ def simulate_game(black_player, white_player):
     )
 
 
-def experience_simulation(num_games, agent1, agent2):
-    collectors = [rl.ExperienceCollector() for _ in range(2)]
+def experience_simulation(num_games, agent1, agent2, progress_callback=None):
+    collectors = [experience.ExperienceCollector() for _ in range(2)]
     agents = [agent1, agent2]
 
     color1 = Player.black
@@ -55,4 +55,7 @@ def experience_simulation(num_games, agent1, agent2):
 
         color1 = color1.other
 
-    return rl.combine_experience(collectors)
+        if progress_callback:
+            progress_callback()
+
+    return experience.combine_experience(collectors)
diff --git a/rl/value.py b/rl/value.py
index 10b22a6..1eac92d 100644
--- a/rl/value.py
+++ b/rl/value.py
@@ -4,7 +4,8 @@ import torch.optim as optim
 from torch.distributions import Categorical
 
 from tugo import encoders
-from tugo import goboard
+# from tugo.game_logic import goboard
+from tugo.game_logic import goboard_fast as goboard
 from tugo.agents import Agent
 from tugo.agents.helpers import is_point_an_eye