diff --git a/pathfinding3d/core/world.py b/pathfinding3d/core/world.py
index e1cb7e6..997d73c 100644
--- a/pathfinding3d/core/world.py
+++ b/pathfinding3d/core/world.py
@@ -62,3 +62,10 @@ def calc_cost(self, node_a: GridNode, node_b: GridNode, weighted: bool = False)
         # TODO: if node_a.grid_id != node_b.grid_id calculate distance between
         # grids as well, for now we ignore switching grids
         return self.grids[node_a.grid_id].calc_cost(node_a, node_b, weighted=weighted)
+
+    def cleanup(self):
+        """
+        Cleanup all grids in this world.
+        """
+        for grid in self.grids.values():
+            grid.cleanup()
diff --git a/pathfinding3d/finder/theta_star.py b/pathfinding3d/finder/theta_star.py
index 8f2e4d6..e44887a 100644
--- a/pathfinding3d/finder/theta_star.py
+++ b/pathfinding3d/finder/theta_star.py
@@ -1,12 +1,12 @@
 import logging
 from typing import Callable, List, Union
 
-from pathfinding3d.core.diagonal_movement import DiagonalMovement
-from pathfinding3d.core.grid import Grid
-from pathfinding3d.core.node import GridNode
-from pathfinding3d.core.util import line_of_sight
-from pathfinding3d.finder.a_star import AStarFinder
-from pathfinding3d.finder.finder import MAX_RUNS, TIME_LIMIT
+from ..core.diagonal_movement import DiagonalMovement
+from ..core.grid import Grid
+from ..core.node import GridNode
+from ..core.util import line_of_sight
+from .a_star import AStarFinder
+from .finder import MAX_RUNS, TIME_LIMIT
 
 
 class ThetaStarFinder(AStarFinder):
@@ -84,19 +84,23 @@ def process_node(
             else than True (used for bi-directional algorithms)
         """
         # Check for line of sight to the grandparent
-        if parent and parent.parent and line_of_sight(grid, node, parent.parent):
-            ng = parent.parent.g + grid.calc_cost(parent.parent, node, self.weighted)
-            if not node.opened or ng < node.g:
-                old_f = node.f
-                node.g = ng
-                node.h = node.h or self.apply_heuristic(node, end)
-                node.f = node.g + node.h
-                node.parent = parent.parent
-                if not node.opened:
-                    open_list.push_node(node)
-                    node.opened = open_value
-                else:
-                    open_list.remove_node(node, old_f)
-                    open_list.push_node(node)
+        if parent and parent.parent and parent.parent.grid_id == node.grid_id:
+            grid_to_use = grid.grids[node.grid_id] if hasattr(grid, "grids") else grid
+            if line_of_sight(grid_to_use, node, parent.parent):
+                ng = parent.parent.g + grid.calc_cost(parent.parent, node, self.weighted)
+                if not node.opened or ng < node.g:
+                    old_f = node.f
+                    node.g = ng
+                    node.h = node.h or self.apply_heuristic(node, end)
+                    node.f = node.g + node.h
+                    node.parent = parent.parent
+                    if not node.opened:
+                        open_list.push_node(node)
+                        node.opened = open_value
+                    else:
+                        open_list.remove_node(node, old_f)
+                        open_list.push_node(node)
+            else:
+                super().process_node(grid, node, parent, end, open_list)
         else:
             super().process_node(grid, node, parent, end, open_list)
diff --git a/test/test_connect_grids.py b/test/test_connect_grids.py
index 2f91e26..52a7fea 100644
--- a/test/test_connect_grids.py
+++ b/test/test_connect_grids.py
@@ -1,6 +1,10 @@
+import pytest
+
 from pathfinding3d.core.grid import Grid
+from pathfinding3d.core.util import expand_path
 from pathfinding3d.core.world import World
 from pathfinding3d.finder.a_star import AStarFinder
+from pathfinding3d.finder.theta_star import ThetaStarFinder
 
 PATH = [
     (2, 0, 0, 0),
@@ -16,8 +20,21 @@
     (2, 0, 0, 1),
 ]
 
+PATH_DIAGONAL = [
+    (2, 0, 0, 0),
+    (2, 1, 0, 0),
+    (2, 1, 1, 0),
+    (2, 2, 2, 0),
+    # move to grid 1
+    (2, 2, 2, 1),
+    (2, 1, 2, 1),
+    (2, 1, 1, 1),
+    (2, 0, 0, 1),
+]
 
-def test_connect():
+
+@pytest.fixture
+def gen_world():
     world0 = [
         [[1, 1, 1], [1, 0, 1], [1, 1, 1]],
         [[1, 1, 1], [1, 0, 1], [1, 1, 1]],
@@ -37,7 +54,32 @@ def test_connect():
 
     # create world with both grids
     world = World({0: grid0, 1: grid1})
+    return world, grid0, grid1
+
+
+def test_connect(gen_world):
+    world, grid0, grid1 = gen_world
+    # clean up the grids
+    world.cleanup()
 
     finder = AStarFinder()
     path, _ = finder.find_path(grid0.node(2, 0, 0), grid1.node(2, 0, 0), world)
     assert [tuple(p) for p in path] == PATH
+
+
+def test_connect_theta_star(gen_world):
+    world, grid0, grid1 = gen_world
+    # clean up the grids
+    world.cleanup()
+
+    finder = ThetaStarFinder()
+    path, _ = finder.find_path(grid0.node(2, 0, 0), grid1.node(2, 0, 0), world)
+    path = [p.identifier for p in path]
+    # path through grid 0 and grid 1
+    paths = {0: [], 1: []}
+    for p in path:
+        paths[p[3]].append(p[:3])
+
+    # expand the paths and combine them with the grid ids
+    path = [(*p, grid_id) for grid_id, path_list in paths.items() for p in expand_path(path_list)]
+    assert path == PATH_DIAGONAL