Collision Detection - Dynamic AABB Tree

This article is for personal use to review the way to create a dynamic AABB tree in Python.

Broad Phase Collision Detection

Before resolving collisions between objects in a simulation, we have to know which object pairs collide. However, when you test collisions between $n$ rigid objects in a physics simulation, there will be a total of $\frac{n(n+1)}{2}$ collision checks, causing the algorithm to work in $O(n^2)$ time. To speed up this, which we will call narrow phase collision detection, we do broad phase collision detection first, which gets possible collision candidate pairs before exact collision detection.

There are many ways for broad phse collision detection, like the sweep and prune method (SAP), spatial subdivision using grids, and bounding volume hierarchy (BVH). SAP is simple but slower compared to others. Spatial subdivision is fast, but usually for particle simulation, which is not the simulation I want since I am dealing with robots. Therefore, I decided to use BVH, which is fast even in environments where gravity easily makes objects cluster. Most existing physics engines like Bullet and PhysicsX use BVH.

Bounding Volume Hierarchy

According to Wikipedia, BVH is described as follows:

A bounding volume hierarchy (BVH) is a tree structure on a set of geometric objects. All geometric objects, that form the leaf nodes of the tree, are wrapped in bounding volumes. These nodes are then grouped as small sets and enclosed within larger bounding volumes. These, in turn, are also grouped and enclosed within other larger bounding volumes in a recursive fashion, eventually resulting in a tree structure with a single bounding volume at the top of the tree.

The above image shows what BVH that uses AABB looks like. Even though there are many bounding volumes, we will use AABB as the bounding volume, not only because it is easy to explain but also because it has the simplest form, which is important for real-time simulation. Like in the image, BVH looks different based on its purpose. For collision detection, we make a tree that has an object's AABB as a leaf. For ray tracking, we make a tree that has a leaf for each part of the object according to criteria.

However, not only is the appearance of those BVHs different, but so is the construction process. For ray tracing, objects usually do not move, so BVH is constructed in such a manner that the tree already has all object parts given. This is called "off-line method". And there are two ways to do that: top-down and bottom-up. However, for collision detection, since objects are dynamic, BVH is constructed with an insertion method where you can add or remove an object from the tree. This is called "online method".

Therefore, in this article, I'll go over the online method, which is the dynamic AABB tree in our case.

Dynamic AABB Tree

My knowledge about the dynamic AABB tree is based on two great articles: James Randall's article and Allen Chou's article. My Python code is the translated version of Chou's C++ code, with slight changes based on Randall's suggestions.

AABB Object

class AABB:  
  def __init__(self, min_x, min_y, min_z, max_x, max_y, max_z, obj=None):
    self.min_x = min_x
    self.min_y = min_y
    self.min_z = min_z
    self.max_x = max_x
    self.max_y = max_y
    self.max_z = max_z
    self.obj = obj
    self.node = None

  def update(self, min_x, min_y, min_z, max_x, max_y, max_z):
    self.min_x = min_x
    self.min_y = min_y
    self.min_z = min_z
    self.max_x = max_x
    self.max_y = max_y
    self.max_z = max_z

  def volume(self):
    return (self.max_x - self.min_x) * (self.max_y - self.min_y) * (self.max_z - self.min_z)

  def union(aabbs):
    min_x = min([obj.min_x for obj in aabbs])
    min_y = min([obj.min_y for obj in aabbs])
    min_z = min([obj.min_z for obj in aabbs])
    max_x = max([obj.max_x for obj in aabbs])
    max_y = max([obj.max_y for obj in aabbs])
    max_z = max([obj.max_z for obj in aabbs])
    return AABB(min_x, min_y, min_z, max_x, max_y, max_z)

  def contains(self, other):
    return (self.min_x <= other.min_x and
            self.max_x >= other.max_x and
            self.min_y <= other.min_y and
            self.max_y >= other.max_y and
            self.min_z <= other.min_z and
            self.max_z >= other.max_z)

  def intersects(self, other):
    return (self.min_x <= other.max_x and
            self.max_x >= other.min_x and
            self.min_y <= other.max_y and
            self.max_y >= other.min_y and
            self.min_z <= other.max_z and
            self.max_z >= other.min_z)

The AABB object is simple. It additionally saves the object's pointer if it is assigned to it and also saves the AABB node's pointer. The reason for saving the object's pointer will be discussed later. The AABB object has several basic functions: The "union" function returns a new AABB that contains multiple AABBs; the "contains" function determines whether some AABB is completely contained within the current AABB; and the "intersects" function determines whether some AABB intersects with the current AABB.

AABB Node

class AABBNode:
  def __init__(self):
    self.parent = []
    self.children = []
    self.collision_check = False
    self.data = None
    self.aabb = None

  def is_leaf(self):
    return not self.children

  def set_branch(self, n0, n1):
    n0.parent, n1.parent = [self], [self]
    self.children = [n0, n1]

  def set_leaf(self, data):
    self.data, data.node = data, self
    self.children = []

  def update_aabb(self):
    if self.is_leaf():
      margin = 0.2
      self.aabb = AABB(self.data.min_x - margin, self.data.min_y - margin,
                       self.data.min_z - margin, self.data.max_x + margin,
                       self.data.max_y + margin, self.data.max_z + margin)
    else:
      self.aabb = AABB.union([child.aabb for child in self.children])

  def get_sibling(self):
    if self == self.parent[0].children[0]:
      return self.parent[0].children[1]
    else:
      return self.parent[0].children[0]

The AABB node is the basic building block for the AABB tree. It can be the leaf node that contains the object's aabb, or it can be the branch node whose aabb contains the child nodes. The interesting part is that "self.data" is the actual AABB of the object and "self.aabb" is the inflated AABB of the object if the node is a leaf of the tree. This will be described later with some explanationary images.

AABB Tree

class AABBTree:
  def __init__(self):
    self.root = []
    self.collision_pairs = []
    self.invalid_nodes = []
    self.tree_updated = True

One change from Chou's code is that I used the tree_updated flag, which tells us if we need to recalculate collision pairs or not. If the tree is updated, we need to recalculate the collision pairs since there has been a change. And I used Python lists as an easy implementation of pointers of C++.

Insert and Insert Node

  def insert(self, aabb):
    self.tree_updated = True
    # if root exists create a new leaf node and insert the node
    if self.root:
      node = AABBNode()
      node.set_leaf(aabb)
      node.update_aabb()
      self.insert_node(node, self.root, index=0)
    # if there is not root make the root as the leaf node
    else:
      self.root = [AABBNode()]
      self.root[0].set_leaf(aabb)
      self.root[0].update_aabb()

  def insert_node(self, node, parent, index=0):
    stack = [(node, parent, index)]
    while stack:
      node, parent, index = stack.pop()
      parent_node = parent[index]

      # if parent node is leaf, split the node and insert children nodes
      if parent_node.is_leaf():
        new_parent = AABBNode()
        new_parent.parent = parent_node.parent
        new_parent.set_branch(node, parent_node)
        parent[index] = new_parent
      # if parent node is branch, compare the children's volume increase
      # and insert the node to the smaller volume increased child
      else:
        aabb0 = parent_node.children[0].aabb
        aabb1 = parent_node.children[1].aabb
        volume_diff0 = AABB.union([aabb0, node.aabb]).volume() - aabb0.volume()
        volume_diff1 = AABB.union([aabb1, node.aabb]).volume() - aabb1.volume()

        if volume_diff0 < volume_diff1:
          stack.append((node,  parent_node.children, 0))
        else:
          stack.append((node,  parent_node.children, 1))

    # update the parent nodes' AABB
    while parent_node:
      parent_node.update_aabb()
      if parent_node.parent:
        parent_node = parent_node.parent[0]
      else: parent_node = None

One important thing about the insert function is that it balances the tree using the cost function, which is the minimum volume increase. The minimum surface increase can also be used, but that is more suitable for BVH for ray tracing. The changes I made are that I have used stack instead of recursive expressions, and I have made the tree's AABBs update when a new node is inserted.

Remove and Remove Node

  def remove(self, aabb):
    # remove the two-way link of the node
    # and remove the node from the tree
    self.tree_updated = True
    node = aabb.node
    node.data = None
    aabb.node = None
    self.remove_node(node)

  def remove_node(self, node):
    if node.parent:
      parent = node.parent[0]
      sibling = node.get_sibling()
      # if node's grandparent exits, delete the parent node
      # and attach the sibling node as the gradparent's child node
      if parent.parent:
        grandparent_node = parent.parent[0]
        sibling.parent[0] = grandparent_node
        if parent == grandparent_node.children[0]:
          grandparent_node.children[0] = sibling
        else: grandparent_node.children[1] = sibling
        # update the grandparent nodes' AABB
        while grandparent_node:
          grandparent_node.update_aabb()
          if grandparent_node.parent:
            grandparent_node = grandparent_node.parent[0]
          else: grandparent_node = None
      # if there is no grandparent, make the sibling node as the root
      else: 
        self.root[0] = sibling
        sibling.parent = []
      node.parent = []
    # if there nothing other than the node, empty the tree
    else:
      self.root = []

Like the insert function, I changed the recursive expression to a while loop and made sure to update AABBs after a node's deletion.

Collision Detection

  def clear_collision_check(self, node):
    # change all cllision checks to False
    stack = [node]
    while stack:
      node = stack.pop()
      node.collision_check = False
      if not node.is_leaf():
        for child_node in node.children:
          stack.append(child_node)

  def child_collision_check(self, node, deque):
    # check node's internal collision if it is not checked
    if not node.collision_check:
      deque.insert(0,(node.children[0], node.children[1]))
      node.collision_check = True

  def compute_pairs(self, node0, node1):
    deque = [(node0, node1)]
    while deque:
      node0, node1 = deque.pop(0)
      if node0.is_leaf():
        # if all nodes are leaf, do AABB collision check
        if node1.is_leaf():
          if node0.aabb.intersects(node1.aabb):
            self.collision_pairs.append([node0.data.obj, node1.data.obj])
        # if node1 is branch, check its internal collision
        else:
          self.child_collision_check(node1, deque)
          deque.extend([(node0, child_node1) for child_node1 in node1.children])
      else:
        # if node0 is branch, check its internal collision
        if node1.is_leaf():
          self.child_collision_check(node0, deque)
          deque.extend([(child_node0, node1) for child_node0 in node0.children])
        # if all nodes are branch, check their internal collisions
        # then check their childrens collisions
        else:
          self.child_collision_check(node0, deque)
          self.child_collision_check(node1, deque)
          deque.extend([(child_node0, child_node1) for child_node0 in node0.children for child_node1 in node1.children])

  def query_collision(self):
    # if tree is updated, compute collision pairs again
    if self.tree_updated:
      self.tree_updated = False
      self.collision_pairs = []
      if self.root[0].is_leaf(): return self.collision_pairs
      self.clear_collision_check(self.root[0])
      self.compute_pairs(self.root[0].children[0], self.root[0].children[1])
    return self.collision_pairs

Chou used "self.collision_check" not to create duplicated collision pairs. However, this made changing the recursive expression to a while loop a little tricky. However, I managed to do it by using the deque concept and putting "child_collision_check" to be checked first before collision checks between the nodes. Another change is that collision pairs get calculated only when there is a change to the tree.

Update Tree

  def update(self):
    if self.root[0]:
      if self.root[0].is_leaf():
        self.root[0].update_aabb()
      else:
        # for nodes that whose AABB no longer valid
        # remove the node and reinsert it        
        self.invalid_nodes = []
        self.valid_check(self.root[0], self.invalid_nodes)

        for node in self.invalid_nodes:
          self.remove_node(node)
          node.update_aabb()
          self.insert_node(node, self.root, index=0)

        if self.invalid_nodes:
          self.tree_updated = True

        self.invalid_nodes = []

  def valid_check(self, node, invalid_nodes):
    # check if actual AABB is out of node's AABB
    stack = [node]
    while stack:
      node = stack.pop()
      if node.is_leaf():
        if not node.aabb.contains(node.data):
          invalid_nodes.append(node)
      else:
        stack.append(node.children[0])
        stack.append(node.children[1])

This part is the art of the dynamic AABB tree. Since objects move a lot in the simulation, the AABB tree has to be rebuilt each time an object moves. However, this is a very costly process for real-time simulation. Therefore, instead of updating every node, we only update nodes whose actual AABB is higher than the buffed AABB of the node. This is why we use the insertion method for building BVH and separate "self.data" and "self.aabb" in the node object.

In the node object code, we simply inflated the AABB by a certain margin. However, if the object is small or moves fast, this is unlikely to be effective. Then, we can also consider the velocity of the object. The updated code will look like this.

  def update_aabb(self, delta_t):
    if self.is_leaf():
      margin = 0.05

      vx = self.data.obj.vx
      vy = self.data.obj.vy
      vz = self.data.obj.vz
      min_x = self.data.min_x
      min_y = self.data.min_y
      min_z = self.data.min_z
      max_x = self.data.max_x
      max_y = self.data.max_y
      max_z = self.data.max_z

      if vx > 0: max_x += vx * delta_t
      else: min_x += vx * delta_t
      if vy > 0: max_y += vy * delta_t
      else: min_y += vy * delta_t
      if vz > 0: max_z += vz * delta_t
      else: min_z += vz * delta_t

      self.aabb = AABB(min_x, min_y, min_z, max_x, max_y, max_z)
    else:
      self.aabb = AABB.union([child.aabb for child in self.children])

The "delta_t" will affect how frequently the collision pairs will be calculated. I still left a small margin because of the rotation of the object. If the object is long and rotating fast, using some margin can be effective.

Results

The above image is the result of dynamic AABB with spheres. The green boxes are AABBs of a tree's branches, and the blue boxes are AABBs of a tree's leaves. Only colliding objects display the actual AABBs, which are highlighted in red. I have exaggerated the margins for visualization. When Object 5, which has a bigger radius than others, moves within the margin, no recalculation happens for collision pairs. However, like in Figures 3 and 4, if the object moves a lot or some objects disappear or appear, recalculation happens.

The speed is pretty fast. It takes approximately 0.0007 seconds to calculate collision detection for six objects. However, due to the limitations of Python, it is still a lot slower than other implementations when we have 30 or 40 objects. Therefore, I strongly encourage you to implement it in C++ and create a Python wrapper if you really want to use it in Python code. But for testing a few objects, it is okay. And actually, for particle-like objects like in my example, it is better to use grids for broad phase collision detection.