Creating robot behaviors with Python generators

Posted 2007-05-02.

Generators are a powerful feature of the Python programming language. In a nutshell, generators let you write a function that behaves like an iterator. The standard approach to programming robot behaviors is based on state machines. However, robotics code is full of special cases, so a complex behavior will typically end up with a lot of bookkeeping cruft. Generators let us simplify the bookkeeping and express the desired behavior in a straightforward manner.

(Idea originally due to Jim Bruce.)

I’ve worked for several years on RoboCup, the international robot soccer competition. Our software is written in a mixture of C++ (for low-level localization and vision algorithms) and Python (for high-level behaviors). Let’s say we want to write a simple goalkeeper for a robot soccer team. Our keeper will be pretty simple; here’s a list of the requirements:

  1. If the ball is far away, stand in place.
  2. If the ball is near by, dive to block it. Dive to the left if the ball is to the left; dive to the right if the ball is to the right.
  3. If we choose a “dive” action, then “stand” on the next frame, nothing will happen. (Well, maybe the robot will twitch briefly....) So when we choose to dive, we need to commit to sending the same dive command for some time (let’s say one second).

The usual approach to robot behavior design relies on hierarchical state machines. Specifically, we might be in a “standing” state while the ball is far away; when the ball becomes close, we enter a “diving” state that persists for one second. Because of requirement 3, this solution will have a few warts: we need to keep track of how much time we’ve spent in the dive state. Every time we add a special case like this, we need to keep some extra state information around. Since robotics code is full of special cases, we tend to end up with a lot of bookkeeping cruft. In contrast, generators will let us clearly express the desired behavior.

On to the state-machine approach. First, we’ll have a class called Features that abstracts the robot’s raw sensor data. For this example, we only care whether the ball is near/far and left/right, so Features will just contain two boolean variables:

class Features(object):
    ballFar = True
    ballOnLeft = True

Next, we make the goalkeeper. The keeper’s behavior is specified by the next() function, which is called thirty times per second by the robot’s main event loop (every time the on-board camera produces a new image). The next() function returns one of three actions: "stand", "diveLeft", or "diveRight", based on the current values of the Features object. For now, let’s pretend that requirement 3 doesn’t exist.

class Goalkeeper(object):
    def __init__(self, features):
        self.features = features

    def next(self):
        features = self.features
        if features.ballFar:
            return 'stand'
        else:
            if features.ballOnLeft:
                return 'diveLeft'
            else:
                return 'diveRight'

That was simple enough. The constructor takes in the Features object; the next() method checks the current Features values and returns the correct action. Now, how about satisfying requirement 3? When we choose to dive, we need to keep track of two things: how long we need to stay in the "dive" state and which direction we dove. We’ll do this by adding a couple of instance variables (self.diveFramesRemaining and self.lastDiveCommand) to the Goalkeeper class. These variables are set when we initiate the dive. At the top of the next() function, we check if self.diveFramesRemaining is positive; if so, we can immediately return self.lastDiveCommand without consulting the Features. Here’s the code:

class Goalkeeper(object):
    def __init__(self, features):
        self.features = features
        self.diveFramesRemaining = 0
        self.lastDiveCommand = None

    def next(self):
        features = self.features
        if self.diveFramesRemaining > 0:
            self.diveFramesRemaining -= 1
            return self.lastDiveCommand
        else:
            if features.ballFar:
                return 'stand'
            else:
                if features.ballOnLeft:
                    command = 'diveLeft'
                else:
                    command = 'diveRight'
                self.lastDiveCommand = command
                self.diveFramesRemaining = 29
                return command

This satisfies all the requirements, but it’s ugly. We’ve added a couple of bookkeeping variables to the Goalkeeper class. Code to properly maintain these variables is sprinkled all over the next() function. Even worse, the structure of the code no longer accurately represents the programmer’s intent: the top-level if-statement depends on the state of the robot rather than the state of the world. The intent of the original next() function is much easier to discern. (In real code, we could use a state-machine class to tidy things up a bit, but the end result would still be ugly when compared to our original next() function.)

With generators, we can preserve the form of the original next() function and keep the bookkeeping only where it’s needed. If you’re not familiar with generators, you can think of them as a special kind of function. The yield keyword is essentially equivalent to return, but the next time the generator is called, execution continues from the point of the last yield, preserving the state of all local variables. With yield, we can use a for loop to “return” the same dive command the next 30 times the function is called! Lines 11-16 of the below code show the magic:

class GoalkeeperWithGenerator(object):
    def __init__(self, features):
        self.features = features

    def behavior(self):
        while True:
            features = self.features
            if features.ballFar:
                yield 'stand'
            else:
                if features.ballOnLeft:
                    command = 'diveLeft'
                else:
                    command = 'diveRight'
                for i in xrange(30):
                    yield command

Here’s a simple driver script that shows how to use our goalkeepers:

import random

f = Features()
g1 = Goalkeeper(f)
g2 = GoalkeeperWithGenerator(f).behavior()

for i in xrange(10000):
    f.ballFar = random.random() > 0.1
    f.ballOnLeft = random.random() < 0.5
    g1action = g1.next()
    g2action = g2.next()
    print "%s\t%s\t%s\t%s" % (
        f.ballFar, f.ballOnLeft, g1action, g2action)
    assert(g1action == g2action)

… and we’re done! I hope you’ll agree that the generator-based keeper is much easier to understand and maintain than the state-machine-based keeper. You can grab the full source code below and take a look for yourself.

#!/usr/bin/env python

class Features(object):
    ballFar = True
    ballOnLeft = True


class Goalkeeper(object):
    def __init__(self, features):
        self.features = features
        self.diveFramesRemaining = 0
        self.lastDiveCommand = None

    def next(self):
        features = self.features
        if self.diveFramesRemaining:
            self.diveFramesRemaining -= 1
            return self.lastDiveCommand
        else:
            if features.ballFar:
                return 'stand'
            else:
                if features.ballOnLeft:
                    command = 'diveLeft'
                else:
                    command = 'diveRight'
                self.lastDiveCommand = command
                self.diveFramesRemaining = 29
                return command


class GoalkeeperWithGenerator(object):
    def __init__(self, features):
        self.features = features

    def behavior(self):
        while True:
            features = self.features
            if features.ballFar:
                yield 'stand'
            else:
                if features.ballOnLeft:
                    command = 'diveLeft'
                else:
                    command = 'diveRight'
                for i in xrange(30):
                    yield command


import random
f = Features()
g1 = Goalkeeper(f)
g2 = GoalkeeperWithGenerator(f).behavior()

for i in xrange(10000):
    f.ballFar = random.random() > 0.1
    f.ballOnLeft = random.random() < 0.5
    g1action = g1.next()
    g2action = g2.next()
    print "%s\t%s\t%s\t%s" % (
        f.ballFar, f.ballOnLeft, g1action, g2action)
    assert(g1action == g2action)